server: serve should be a context manager

This commit is contained in:
Roman Zeyde
2015-06-17 16:52:11 +03:00
parent 9755c38eb3
commit a60a627b8a
2 changed files with 33 additions and 39 deletions

View File

@@ -68,7 +68,23 @@ def spawn(func, **kwargs):
t.join()
def run(command, environ):
@contextlib.contextmanager
def serve(key_files, signer, sock_path=None):
if sock_path is None:
sock_path = tempfile.mktemp(prefix='ssh-agent-')
keys = [formats.parse_public_key(k) for k in key_files]
environ = {'SSH_AUTH_SOCK': sock_path, 'SSH_AGENT_PID': str(os.getpid())}
with unix_domain_socket_server(sock_path) as server:
with spawn(server_thread, server=server, keys=keys, signer=signer):
try:
yield environ
finally:
log.debug('closing server')
server.shutdown(socket.SHUT_RD)
def run_process(command, environ):
log.debug('running %r with %r', command, environ)
env = dict(os.environ)
env.update(environ)
@@ -80,19 +96,3 @@ def run(command, environ):
ret = p.wait()
log.debug('subprocess %d exited: %d', p.pid, ret)
return ret
def serve(key_files, command, signer, sock_path=None):
if sock_path is None:
sock_path = tempfile.mktemp(prefix='ssh-agent-')
keys = [formats.parse_public_key(k) for k in key_files]
environ = {'SSH_AUTH_SOCK': sock_path, 'SSH_AGENT_PID': str(os.getpid())}
with unix_domain_socket_server(sock_path) as server:
with spawn(server_thread, server=server, keys=keys, signer=signer):
try:
ret = run(command=command, environ=environ)
finally:
log.debug('closing server')
server.shutdown(socket.SHUT_RD)
return ret