import re
import os
import json
import subprocess
import base64
import logging

import boto3

def list_secrets():
    """
    List Secrets Manager secrets with specific pattern.

    Args:
    prefix (str): Optional. Prefix filter for the secret name.

    Returns:
    list: list of Secrets Manager secret names
    """
    client = boto3.client('secretsmanager', region_name='us-east-1')

    env = os.environ.get('ENVIRONMENT', 'qa')
    prefix = f'{env}/direct_delivery/connection_info'

    filters = []
    if prefix:
        filters.append({'Key': 'name', 'Values': [prefix]})

    paginator = client.get_paginator('list_secrets')
    response_iterator = paginator.paginate(Filters=filters, IncludePlannedDeletion=False)

    secret_list = []
    for response in response_iterator:
        if 'SecretList' in response:
            for secret in response['SecretList']:
                secret_list.append(secret['Name'])

    key = rf'{env}/direct_delivery/connection_info.*/release'
    delivery_worker_keys = [secret for secret in secret_list if re.match(key, secret)]

    return delivery_worker_keys

def get_md5_fingerprint(ssh_key):
    """
    Compute the MD5 fingerprint of a given SSH key using the ssh-keygen tool.

    Parameters:
    ssh_key (str): The SSH key string for which the fingerprint needs to be computed.

    Returns:
    str: The MD5 fingerprint of the provided SSH key.
    """
    cmd = ["ssh-keygen", "-l", "-E", "md5", "-f", "-"]
    result = subprocess.run(cmd, capture_output=True, input=ssh_key.encode(), text=False, timeout=10)

    if result.returncode != 0:
        logging.error(f"ssh-keygen failed with error: {result.stderr.decode()}")
        return None

    # Extract the fingerprint
    fingerprint = result.stdout.strip().split()[1][4:].decode('utf-8')
    return fingerprint

def add_known_host(arn, secret_data, value):
    """
    Add or update the 'known_host' key for a given secret in AWS Secrets Manager.

    Parameters:
    secret_data (dict): The secret data dictionary containing the secret's details.
    value (str): The value to set for the 'known_host' key.

    Note:
    This function will overwrite the 'known_host' value if it already exists in the secret.
    """
    session = boto3.session.Session()
    client = session.client(service_name='secretsmanager')

    # Add the new key-value pair (overwrites if key already exists)
    secret_data['known_host'] = value

    # Store the updated secret value back to AWS Secrets Manager
    client.put_secret_value(SecretId=arn, SecretString=json.dumps(secret_data))
    
def main():
    """Add Public Host Keys to Secrets for Delivery Workers."""
    client = boto3.client('secretsmanager')
    delivery_worker_secrets = list_secrets()

    for secret in delivery_worker_secrets:
        secret_value = client.get_secret_value(SecretId=secret)
        secret_data = json.loads(secret_value['SecretString'])

        if secret_data['connection_type'] == "sftp":
            logging.info(f'Connection type for the secret {secret} is sftp')
            domain = secret_data['domain_name']
            port = secret_data['port']

            if domain and port is not None:
                cmd = ["ssh-keyscan", "-t", "rsa", "-p", str(port), domain]

                try:
                    result = subprocess.run(cmd, capture_output=True, text=True, timeout=60, check=True)
                except subprocess.CalledProcessError as e:
                    logging.error(f"Error executing command {cmd}. Error: {e}")
                    continue

                key_output = result.stdout
                print(secret)
                md5_fingerprint = get_md5_fingerprint(key_output)
                if md5_fingerprint is None:
                    logging.warning(f'Skipping secret: {secret} further processing for this key due to ssh-keygen error.')
                    continue
                if md5_fingerprint == secret_data['md5_fingerprint']:
                    logging.info(f'The MD5 fingerprint for the secret {secret} matches.')
                    key_split = result.stdout.split()
                    public_key = f'{key_split[1].replace(" ", "")} {key_split[2]}'
                    encoded = base64.b64encode(public_key.encode()).decode()
                    logging.info(f'Adding public host keys to secrets for delivery workers.')
                    add_known_host(secret_value['ARN'], secret_data, encoded)
                    logging.info(f'The secret {secret} was updated')
                else:
                    logging.warning(f"The MD5 fingerprint for the secret {secret} does not match. Expected: {secret_data['md5_fingerprint']}, Found: {md5_fingerprint}")

if __name__ == '__main__':
    main()