server: serve should be a context manager

sshagent/trezor_agent.py

@@ -22,30 +22,24 @@ def main():
     level = verbosity[min(args.verbose, len(verbosity) - 1)]
     logging.basicConfig(level=level, format=fmt)
 
-    client = trezor.Client(factory=trezor.TrezorLibrary)
+    with trezor.Client(factory=trezor.TrezorLibrary) as client:
+        key_files = []
+        for label in args.labels:
+            pubkey = client.get_public_key(label=label)
+            key_file = formats.export_public_key(pubkey=pubkey, label=label)
+            key_files.append(key_file)
 
-    key_files = []
-    for label in args.labels:
-        pubkey = client.get_public_key(label=label)
-        key_file = formats.export_public_key(pubkey=pubkey, label=label)
-        key_files.append(key_file)
+        if not args.command:
+            sys.stdout.write(''.join(key_files))
+            return
 
-    if not args.command:
-        sys.stdout.write(''.join(key_files))
-        return
+        signer = client.sign_ssh_challenge
 
-    signer = client.sign_ssh_challenge
-
-    ret = -1
-    try:
-        ret = server.serve(
-            key_files=key_files,
-            command=args.command,
-            signer=signer)
-        log.info('exitcode: %d', ret)
-    except KeyboardInterrupt:
-        log.info('server stopped')
-    sys.exit(ret)
+        try:
+            with server.serve(key_files=key_files, signer=signer) as env:
+                return server.run_process(command=args.command, environ=env)
+        except KeyboardInterrupt:
+            log.info('server stopped')
 
 if __name__ == '__main__':
-    main()
+    sys.exit(main())