A script to run a command within an ssh-agent session.
#! /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) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment