ssh-agent.py wrapper

ssh-agent.py

#!/usr/bin/env python
"""\
ssh agent wrapper to detect or run agent, set environment
and either show it or execute specified command in this
environment.

usage: ssh-agent.py
       ssh-agent.py <command>

examples:
    ssh-agent.py ssh-add
    ssh-agent.py ssh -A user@host
"""

"""
[x] no agent is running
  [x] run agent

[x] agent is running
  [x] set environment
  [x] show how to set environment manually

  [x] if there is command, run it
"""

__version__ = '1.0'
__author__  = 'anatoly techtonik <[email protected]>'
__license__ = 'Public Domain'

import os
import sys

#--[inline shellrun 2.0 import run, run_capture]
import subprocess

class Result(object):
    def __init__(self, command=None, retcode=None, output=None):
        self.command = command or ''
        self.retcode = retcode
        self.output = output
        self.success = False
        if retcode == 0:
            self.success = True

def run(command):
    process = subprocess.Popen(command, shell=True)
    process.communicate()
    return Result(command=command, retcode=process.returncode)

def run_capture(command):
    outpipe = subprocess.PIPE
    process = subprocess.Popen(command, shell=True, stdout=outpipe,
                                                    stderr=outpipe)
    output, _ = process.communicate()
    return Result(command=command, retcode=process.returncode, output=output)
#--[/inline]


SSH_ENV = os.path.expanduser("~/.ssh/agent.env")


def detected():
    if 'SSH_AUTH_SOCK' not in os.environ:
        if os.path.exists(SSH_ENV):
            agentenv = parse()
            os.environ.update(agentenv)
        if 'SSH_AUTH_SOCK' not in os.environ:
            return False
    r = run_capture('ssh-add -L')
    if r.retcode == 0:    # has keys
        return True 
    elif r.retcode == 1:  # no keys
        return True
    elif r.retcode == 2:  # not running
        return False
    else:
        print(r.output)
        return False

def start():
    run('ssh-agent > "%s"' % SSH_ENV)

def parse():
    params = open(SSH_ENV, 'r').read()
    # SSH_AUTH_SOCK=/tmp/ssh-owYJTz7968/agent.7968; export SSH_AUTH_SOCK;
    # SSH_AGENT_PID=6284; export SSH_AGENT_PID;
    # echo Agent pid 6284;
    params = params.replace('; ', '\n')
    params = params.replace(';', '')
    result = dict()
    for line in params.splitlines():
        if '=' not in line:
            continue
        key, value = line.split('=')
        result[key] = value
    return result


def main():
    if '-h' in sys.argv or '--help' in sys.argv:
        sys.exit(__doc__)

    if not detected():
        print('..agent not detected, starting..')
        start()

    if not detected:
        sys.exit('error starting ssh-agent')
    else:
        if os.name == 'nt':
            print('@rem ..agent is running, set environment as:')
            for k,v in parse().items():
                print('set %s=%s' % (k,v))
        else:
            print('# ..agent is running, set environment as:')
            run('cat "%s"' % SSH_ENV)

        if sys.argv[1:]:
            run(sys.argv[1:])


if __name__ == '__main__':
    main()