Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
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
You can’t perform that action at this time.