Created
March 14, 2017 20:31
-
-
Save toolness/c75a3b5cef01ffc1bb2d9ba89c8da63e to your computer and use it in GitHub Desktop.
A script to run a command within an ssh-agent session.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#! /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
thanks a lot @toolness ! this helped me for signing git commits with an SSH key because only ssh-agent is supported.