scripts/q7s-cp.py

@@ -11,16 +11,27 @@ def main():
     print(f"Running command: {cmd}")
     result = os.system(cmd)
     if result != 0:
-        print("")
+        prompt_ssh_key_removal()
-        print("Removing problematic SSH key and trying again..")
-        remove_ssh_key_cmd = (
-            'ssh-keygen -f "${HOME}/.ssh/known_hosts" -R "[localhost]:1535"'
-        )
-        os.system(remove_ssh_key_cmd)
         print(f'Running command "{cmd}"')
         result = os.system(cmd)
 
 
+def prompt_ssh_key_removal():
+    do_remove_key = input("Do you want to remove problematic keys on localhost ([Y]/n)?: ")
+    if not do_remove_key.lower() in ["y", "yes", "1", ""]:
+        sys.exit(1)
+    port = 0
+    while True:
+        port = input("Enter port to remove: ")
+        if not port.isdecimal():
+            print("Invalid port detected")
+        else:
+            break
+    cmd = f'ssh-keygen -f "$ {{HOME}}/.ssh/known_hosts" -R "[localhost]:{port}"'
+    print(f"Removing problematic SSH key with command {cmd}..")
+    os.system(cmd)
+
+
 def handle_args():
     help_string = (
         "This script copies files to the Q7S as long as port forwarding is active.\n"
@@ -58,8 +69,7 @@ def handle_args():
     )
     # Positional argument(s)
     parser.add_argument(
-        "source", help="Source files to copy or target files to copy back to host",
+        "source", help="Source files to copy or target files to copy back to host"
-         nargs="+"
     )
     return parser.parse_args()
 
@@ -72,10 +82,6 @@ def build_cmd(args):
     address = ""
     port_args = ""
     target = args.target
-    if args.invert and len(args.source) > 1:
-        print("Multiple source files not allowed for inverse copying")
-        sys.exit(1)
-    source_files = " ".join(args.source)
     if args.flatsat:
         address = "eive@flatsat.eive.absatvirt.lw"
     else:
@@ -92,9 +98,9 @@ def build_cmd(args):
         else:
             target = args.target
     if args.invert:
-        cmd += f"{port_args} {address}:{source_files} {target}"
+        cmd += f"{port_args} {address}:{args.source} {target}"
     else:
-        cmd += f"{port_args} {source_files} {address}:{target}"
+        cmd += f"{port_args} {args.source} {address}:{target}"
     return cmd