#!/usr/bin/python

'''
Documentation:
---
module: ssh_agent
short_description: Add SSH private key
version_added: "0.0.12"
description:
  - This SSH module allows adding private keys into the authentication agent.
options:
  ssh_key_file:
    description:
      - The path to the private key file to add to the SSH agent.
    required: true
  ssh_env_file:
    description:
      - The path to the SSH environment file.
    required: true
  ssh_passphrase:
    description:
      - The passphrase for the private key (if applicable).
    required: false
'''

import os
import sys
import subprocess
import configparser

try:
    from ansible.module_utils.six.moves import configparser
except ImportError:
    from ansible.module_utils._text import configparser

from ansible.module_utils.basic import AnsibleModule

def start_agent():
    """
    Start a new SSH agent and write environment variables to file
    """
    cmd = subprocess.Popen(['ssh-agent', '-s'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    stdout, stderr = cmd.communicate()
    if cmd.returncode != 0:
        raise Exception('Error starting ssh-agent: {}'.format(stderr.decode()))
    # Parse SSH agent output to get environment variables
    env_vars = dict(line.split('=') for line in stdout.decode().split('\n') if '=' in line)
    # Write environment variables to file in INI format
    config = configparser.ConfigParser()
    config['default'] = env_vars
    with open(os.path.expanduser('~/.ssh-agent-env'), 'w') as f:
        config.write(f)

def main():
    module_args = dict(
        ssh_key_file=dict(type='str', required=True),
        ssh_env_file=dict(required=True, type='str'),
        ssh_passphrase=dict(type='str', required=False, default=None, no_log=True)
    )
    module = AnsibleModule(
        argument_spec=module_args,
        supports_check_mode=False
    )

    key = module.params['ssh_key_file']
    env = module.params['ssh_env_file']
    passphrase = module.params['ssh_passphrase']

    env_vars = dict(os.environ)
    if os.path.exists(env):
        config = configparser.ConfigParser()
        config.read(env)
        if 'SSH_AUTH_SOCK' in config['default']:
            env_vars['SSH_AUTH_SOCK'] = config['default']['SSH_AUTH_SOCK']
            if passphrase is not None:
                cmd = subprocess.Popen(['ssh-add', '{key}'], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env_vars)
                stdout, stderr = cmd.communicate(input=passphrase.encode())
                if cmd.returncode != 0:
                    module.fail_json(msg='Error adding key and passphrase to agent: {}'.format(stderr.decode()))
            else:
                cmd = subprocess.Popen(['ssh-add', '{key}'], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env_vars)
                stdout, stderr = cmd.communicate()
                if cmd.returncode != 0:
                    module.fail_json(msg='Error adding key to agent: {}'.format(stderr.decode()))
            module.exit_json(changed=True, msg='Added key to agent')
        else:
            module.fail_json(msg='SSH_AUTH_SOCK not found in %s' % env)
    else:
        module.fail_json(msg='%s does not exist' % env)

if __name__ == '__main__':
    start_agent()
    main()