diff --git a/src/MigrateUpgradeTrait.php b/src/MigrateUpgradeTrait.php
index f857933..914f3f7 100644
--- a/src/MigrateUpgradeTrait.php
+++ b/src/MigrateUpgradeTrait.php
@@ -50,13 +50,10 @@ trait MigrateUpgradeTrait {
         ->set('source.database', $database);
 
       // Configure file migrations so they can find the files.
-      // @todo: Handle D7.
-      if ($migration_id === 'd6_file' || $migration_id === 'd6_user_picture_file') {
-        if ($site_address) {
-          // Make sure we have a single trailing slash.
-          $site_address = rtrim($site_address, '/') . '/';
-          $config->set('destination.source_base_path', $site_address);
-        }
+      if ($config->get('destination.plugin') == 'entity:file' && $site_address) {
+        // Make sure we have a single trailing slash.
+        $site_address = rtrim($site_address, '/') . '/';
+        $config->set('destination.source_base_path', $site_address);
       }
       $config->save();
     }