A script to run a command within an ssh-agent session.

sshify.py

#! /usr/bin/env python3

"""
    This utility will execute the given command (by default, your shell)
    in a subshell, with an ssh-agent process running and your
    private key added to it. When the subshell exits, the ssh-agent
    process is killed.
"""

# This code was written by Atul Varma in February 2017. It requires
# Python 3.6.
#
# License: CC0 1.0 Universal (Public Domain)

import argparse
import os
import re
import subprocess
import sys
from typing import Dict, List
import unittest


class SshAgent(object):
    def __init__(self, agent_env: Dict[str, str]) -> None:
        self.agent_env = agent_env
        self.pid = agent_env['SSH_AGENT_PID']
        self.environ = {}  # type: Dict[str, str]
        self.environ.update(os.environ)
        self.environ.update(agent_env)
        self.environ.update({
            'PROMPT': '(sshify) ' + self.environ.get('PROMPT', ''),
        })

    @classmethod
    def start(cls) -> 'SshAgent':
        print("Starting ssh-agent")
        output = subprocess.check_output(['ssh-agent', '-s'])
        agent_env = cls.parse_agent_env(output)
        return cls(agent_env)

    def add_key(self, path: str) -> None:
        subprocess.check_call(['ssh-add', path], env=self.environ)

    def stop(self) -> None:
        print("Killing ssh-agent")
        subprocess.check_call(['kill', self.pid])

    def __enter__(self) -> 'SshAgent':
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.stop()

    @staticmethod
    def parse_agent_env(output: bytes) -> Dict[str, str]:
        result = {}
        for name, value in re.findall(r'([A-Z_]+)=([^;]+);',
                                      output.decode('ascii')):
            result[name] = value
        return result


def main(argv: List[str]) -> None:
    parser = argparse.ArgumentParser(
        description=__doc__,
        prog=os.path.basename(argv[0]),
        usage='%(prog)s [options] [command]',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("--test", help="Run test suite and exit",
                        action="store_true")
    parser.add_argument("--key", help="Path to SSH key",
                        default="~/.ssh/id_rsa")
    args, cmd = parser.parse_known_args(argv[1:])

    if args.test:
        unittest.main(argv=[argv[0]] + cmd)
        return

    if len(cmd) == 0:
        cmd = [os.environ.get('COMSPEC', os.environ.get('SHELL', 'bash'))]

    env = {}  # type: Dict[str, str]

    def run():
        try:
            subprocess.call(cmd, env=env)
        except KeyboardInterrupt:
            pass

    if 'SSH_AUTH_SOCK' not in os.environ:
        with SshAgent.start() as agent:
            agent.add_key(args.key)
            env.update(agent.environ)
            run()
    else:
        # We've already got ssh-agent running, just run the command.
        env.update(os.environ)
        run()


class Tests(unittest.TestCase):
    EXAMPLE_OUTPUT = b"""\
    SSH_AUTH_SOCK=/tmp/ssh-H1qkEfCM8owV/agent.8420; export SSH_AUTH_SOCK;
    SSH_AGENT_PID=12588; export SSH_AGENT_PID;
    echo Agent pid 12588;
    """

    def test_parsing_works(self):
        self.assertEqual(
            SshAgent.parse_agent_env(self.EXAMPLE_OUTPUT),
            {
                'SSH_AUTH_SOCK': '/tmp/ssh-H1qkEfCM8owV/agent.8420',
                'SSH_AGENT_PID': '12588',
            }
        )

    def test_mypy(self):
        try:
            import mypy  # type: ignore
            subprocess.check_call([sys.executable, '-m', 'mypy', __file__])
        except ModuleNotFoundError:
            raise unittest.SkipTest()


if __name__ == "__main__":
    main(sys.argv)