run.py

import argparse
import os
import re
import sys
from datetime import datetime

import paramiko
from dotenv import load_dotenv


def establishConnection(ssh_server, ssh_port, user, pasd, public_key_auth, pk_file_location):
    ssh_client = paramiko.client.SSHClient()
    ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

    try:
        if public_key_auth:
            print('asfd')
            ssh_client.connect(
                hostname=ssh_server,
                port=ssh_port,
                username=user,
                key_filename=pk_file_location
            )
        else:
            ssh_client.connect(hostname=ssh_server,
                            port=ssh_port,
                            username=user,
                            password=pasd
            )
    except paramiko.AuthenticationException:
        print("Authentication failed. Please check your username and password.")
        ssh_client.close()
        return None
    except paramiko.SSHException as e:
        print(f"SSH error: {e}")
        ssh_client.close()
        return None

    return ssh_client

def isCommandInstalled(base_cmd, ssh_client):
    full_command = f"which {base_cmd}"
    _stdin, stdout, _stderr = ssh_client.exec_command(full_command)
    result = re.match(stdout.read().decode('utf-8'), 'not\ found')
    return result

def runCommand(base_cmd, arguments, save_out, no_out):
    try:
        if save_out:
            output_file_name = datetime.now().strftime(f"%d-%m-%Y_%H-%M-%S_{base_cmd}")
            makedir = f"mkdir -p $HOME/offloaded_scans/ && "
            out_redirect = f"| tee $HOME/offloaded_scans/{output_file_name}.txt"
            full_command = f"{makedir} {base_cmd} {arguments} {out_redirect}"
        else:
            full_command = f"{base_cmd} {arguments}"
        _stdin, stdout, _stderr = ssh_client.exec_command(full_command)
        output = stdout.read().decode('utf-8')
        if no_out == False:
            print(output)
        if save_out:
            print(f"File saved on remote machine under $HOME/offloaded_scans/ with name {output_file_name}.txt", end="\n\n")
    except paramiko.AuthenticationException:
        print("Authentication failed. Please check your username and password.")
    except paramiko.SSHException as e:
        print(f"SSH error: {e}")
    finally:
        ssh_client.close()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Python script to offload resource intesive tasks to IoT devices"
    )

    required_args = parser.add_argument_group("required arguments")
    optional_args = parser.add_argument_group("optional arguments")

    required_args.add_argument(
        "-b",
        "--base",
        dest="base",
        help="Base command",
        required=True
    )
    required_args.add_argument(
        "-a",
        "--args",
        dest="args",
        help="Arguments for the base command",
        required=True
    )

    optional_args.add_argument(
        "-s",
        "--save-output-remote",
        dest="sor",
        help="Save to file on remote device",
        action=argparse.BooleanOptionalAction,
        default=False
    )
    optional_args.add_argument(
        "-n",
        "--no-output-local",
        dest="nol",
        help="Do not show output to in the terminal and exit as soon as command has been supplied",
        action=argparse.BooleanOptionalAction,
        default=False
    )
    optional_args.add_argument(
        "-p",
        "--use-public-key-remote",
        dest="upkr",
        help="Use public key authentication for ssh connection",
        action=argparse.BooleanOptionalAction,
        default=False
    )

    try:
        args = parser.parse_args()
    except:
        parser.print_help()
        sys.exit(1)

    base_cmd = args.base
    arguments = args.args

    save_out = args.sor
    no_out = args.nol
    public_key_auth = args.upkr

    load_dotenv()

    server = os.environ.get("server")
    port = os.environ.get("port")
    username = os.environ.get("user")
    password = os.environ.get("password")
    pk_file_location = ""

    if public_key_auth:
        pk_file_location = os.environ.get("ssh_public_key")

    ssh_client = establishConnection(
        server,
        port,
        username,
        password,
        public_key_auth,
        pk_file_location
    )

    if ssh_client != None:
        check = isCommandInstalled(
            base_cmd,
            ssh_client
        )

        if check == None:
            runCommand(
                base_cmd,
                arguments,
                save_out,
                no_out
            )
        else:
            print(f"{base_cmd} has not been installed on the server")