eive-obsw/scripts/q7s-cp.py

#!/usr/bin/env python3
import argparse
import os
import sys


def main():
    args = handle_args()
    cmd = build_cmd(args)
    # Run the command
    print(f"Running command: {cmd}")
    result = os.system(cmd)
    if result != 0:
        print("")
        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 handle_args():
    help_string = (
        "This script copies files to the Q7S as long as port forwarding is active.\n"
    )
    help_string += (
        "You can set up port forwarding with "
        '"ssh -L 1535:192.168.133.10:22 <eive-flatsat-ip>" -t /bin/bash'
    )
    parser = argparse.ArgumentParser(description=help_string)
    # Optional arguments
    parser.add_argument(
        "-r", "--recursive", dest="recursive", default=False, action="store_true"
    )
    parser.add_argument(
        "-t",
        "--target",
        help="Target destination. If files are copied to Q7S, will be /tmp by default. "
        "If files are copied back to host, will be current directory by default",
        default="",
    )
    parser.add_argument("-P", "--port", help="Target port", default=1535)
    parser.add_argument(
        "-i",
        "--invert",
        default=False,
        action="store_true",
        help="Copy from Q7S to host instead. Always copies to current directory.",
    )
    parser.add_argument(
        "-f",
        "--flatsat",
        default=False,
        action="store_true",
        help="Copy to flatsat instead"
    )
    # Positional argument(s)
    parser.add_argument(
        "source", help="Source files to copy or target files to copy back to host",
         nargs="+"
    )
    return parser.parse_args()


def build_cmd(args):
    # Build run command
    cmd = "scp "
    if args.recursive:
        cmd += "-r "
    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:
        address = "root@localhost"
        port_args=f"-P {args.port}"
    if args.invert:
        if target == "":
            target = "."
        else:
            target = args.target
    else:
        if target == "":
            target = f"/tmp"
        else:
            target = args.target
    if args.invert:
        cmd += f"{port_args} {address}:{source_files} {target}"
    else:
        cmd += f"{port_args} {source_files} {address}:{target}"
    return cmd


if __name__ == "__main__":
    main()