eive-obsw/scripts/q7s-cp.py

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


DEFAULT_PORT = 1539


def main():
    args = handle_args()
    cmd = build_cmd(args)
    # Run the command
    print(f"Running command: {cmd}")
    result = os.system(cmd)
    if result != 0:
        prompt_ssh_key_removal()
        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 do_remove_key.lower() not 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"
    )
    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=DEFAULT_PORT)
    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"
    )
    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.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
    # accepted_key_rsa_args = "-o HostKeyAlgorithms=+ssh-rsa -o PubkeyAcceptedKeyTypes=+ssh-rsa"
    accepted_key_rsa_args = ""
    if args.invert:
        cmd += f"{port_args} {accepted_key_rsa_args} {address}:{args.source} {target}"
    else:
        cmd += f"{port_args} {accepted_key_rsa_args} {args.source} {address}:{target}"
    return cmd


if __name__ == "__main__":
    main()