#!/usr/bin/env python3

## This script uses a CSV file with the following columns
## to automate provisioning student EC2 instances with AWS
## by running a script. This script is meant is not meant to
## be ran by the user (e.g. students). It also requires an SSH
## keyfile (which is assumed to be shared between instances)
## to log into the instances and run the SSH commands.
#
# CSV columns:
#
# 1. EC2 instance public IP address (public_ip_address)
# 2. GitHub username (github_username)
#
# Usage:
#
# ./setup-instances.py <path_to_csv> <path_to_keyfile> --script-url <script_url>
#
# Example:
#
# ./setup-instances.py ./users.csv ./openchami.pem https://gist.githubusercontent.com/alexlovelltroy/96bfc8bb6f59c0845617a0dc659871de/raw
from fabric import Connection
from argparse import ArgumentParser
from csv import DictReader

DEFAULT_SCRIPT_URL = "https://gist.githubusercontent.com/davidallendj/bb57c25adefb3e9921ccc3502e9ca13b/raw/2165263c07edc59c007b015c8d21f245f403c09d/setup-student.sh"

def read_user_data(csv_path: str, ignore_header: bool) -> list:
    data = []
    with open(csv_path, 'r') as file:
        if ignore_header:
            next(file)
        reader = DictReader(file)
        for row in reader:
            data.append(row)
    return data


def ssh_setup_instance(
    remote_host: str,
    username: str,
    keyfile_path: str,
    script_url: str = DEFAULT_SCRIPT_URL
) -> None:
    # download the script on the remote host instance
    conn = Connection(
        host=f"rocky@{remote_host}",
        connect_kwargs={"key_filename": keyfile_path}
    )
    result = conn.run(f"curl -fssL {script_url} -o ./setup-student.sh && chmod 755 ./setup-student.sh && ./setup-student.sh {username}", hide=True)
    print("{0.stdout}\n".format(result))


def main() -> None:
    # setup command-line arguments
    parser = ArgumentParser(
        prog='setup-instance',
        description='Automates setting up an EC2 instance for the OpenCHAMI tutorial.',
        epilog='./setup-instances.py <path_to_csv> <path_to_keyfile> <setup_script_url>'
    )
    parser.add_argument('csv-file', help="Path to CSV containing user information.", default="")
    parser.add_argument('keyfile', help="Path to key file used to log into AWS EC2 instance. This is assumed to be shared between instances.")
    parser.add_argument('--script-url', dest="script-url", help="URL pointing to setup script to initialize EC2 instance.", default=DEFAULT_SCRIPT_URL)
    parser.add_argument('--ignore-header', dest="ignore-header", type=bool, help="Ignores the first line in CSV file.", default=False)
    args = vars(parser.parse_args())


    # read contents from CSV file and load into dict
    data: list = read_user_data(args["csv-file"], args["ignore-header"])
    if len(data) == 0:
        print("no data found in CSV file")
        exit()

    # run SSH commands to setup instance
    for row in data:
        print(f"{row}")
        ip_addr: str = row["public_ip_address"]
        username: str = row["github_username"]
        if not ip_addr or not username:
            continue
        ssh_setup_instance(ip_addr, username, args["keyfile"], args["script-url"])


if __name__ == "__main__":
    main()