From eb525e1b622437eec6358b4a614defb7d4e50e6e Mon Sep 17 00:00:00 2001
From: Roman Zeyde <roman.zeyde@gmail.com>
Date: Wed, 26 Apr 2017 22:05:49 +0300
Subject: [PATCH] gpg: simplify Python entry point and refactor Bash scripts a
 bit

Now there is a single 'trezor-gpg' tool, with various subcommands.
---
 scripts/gpg-agent            |  2 ++
 scripts/gpg-init             |  4 +--
 scripts/gpg-shell            |  4 +--
 setup.py                     |  4 +--
 trezor_agent/gpg/__main__.py | 56 +++++++++++++++++++++++-------------
 5 files changed, 43 insertions(+), 27 deletions(-)
 create mode 100755 scripts/gpg-agent

diff --git a/scripts/gpg-agent b/scripts/gpg-agent
new file mode 100755
index 0000000..0ea74e7
--- /dev/null
+++ b/scripts/gpg-agent
@@ -0,0 +1,2 @@
+#!/bin/bash
+trezor-gpg agent
diff --git a/scripts/gpg-init b/scripts/gpg-init
index 3c4cdba..83496d2 100755
--- a/scripts/gpg-init
+++ b/scripts/gpg-init
@@ -14,7 +14,7 @@ mkdir -p "${HOMEDIR}"
 chmod 700 "${HOMEDIR}"
 
 # Generate new GPG identity and import into GPG keyring
-trezor-gpg-create -v "${USER_ID}" -t "${TIMESTAMP}" -e "${CURVE}" > "${HOMEDIR}/pubkey.asc"
+trezor-gpg create -v "${USER_ID}" -t "${TIMESTAMP}" -e "${CURVE}" > "${HOMEDIR}/pubkey.asc"
 gpg2 --homedir "${HOMEDIR}" --import < "${HOMEDIR}/pubkey.asc"
 rm -f "${HOMEDIR}/S.gpg-agent"  # (otherwise, our agent won't be started automatically)
 
@@ -24,7 +24,7 @@ echo "${FINGERPRINT}:6" | gpg2 --homedir "${HOMEDIR}" --import-ownertrust
 
 # Prepare GPG configuration file
 echo "# TREZOR-based GPG configuration
-agent-program $(which trezor-gpg-agent)
+agent-program $(dirname ${0})/gpg-agent
 personal-digest-preferences SHA512
 " | tee "${HOMEDIR}/gpg.conf"
 
diff --git a/scripts/gpg-shell b/scripts/gpg-shell
index fbbf046..b4db164 100755
--- a/scripts/gpg-shell
+++ b/scripts/gpg-shell
@@ -13,10 +13,10 @@ then
 fi
 
 # Make sure that the device is unlocked before starting the shell
-trezor-gpg-unlock
+trezor-gpg unlock
 
 # Make sure TREZOR-based gpg-agent is running
-gpg-connect-agent --agent-program "$(which trezor-gpg-agent)" </dev/null
+gpg-connect-agent --agent-program "$(dirname $0)/gpg-agent" </dev/null
 
 COMMAND=$*
 if [ -z "${COMMAND}" ]
diff --git a/setup.py b/setup.py
index 6435b83..ee526f8 100644
--- a/setup.py
+++ b/setup.py
@@ -34,8 +34,6 @@ setup(
     ],
     entry_points={'console_scripts': [
         'trezor-agent = trezor_agent.__main__:run_agent',
-        'trezor-gpg-create = trezor_agent.gpg.__main__:main_create',
-        'trezor-gpg-agent = trezor_agent.gpg.__main__:main_agent',
-        'trezor-gpg-unlock = trezor_agent.gpg.__main__:auto_unlock',
+        'trezor-gpg = trezor_agent.gpg.__main__:main',
     ]},
 )
diff --git a/trezor_agent/gpg/__main__.py b/trezor_agent/gpg/__main__.py
index 05671ad..e1cceb0 100755
--- a/trezor_agent/gpg/__main__.py
+++ b/trezor_agent/gpg/__main__.py
@@ -15,7 +15,7 @@ from .. import device, formats, server, util
 log = logging.getLogger(__name__)
 
 
-def run_create(args):
+def export_public_key(args):
     """Generate a new pubkey for a new/existing GPG identity."""
     log.warning('NOTE: in order to re-generate the exact same GPG key later, '
                 'run this command with "--time=%d" commandline flag (to set '
@@ -64,16 +64,8 @@ def run_create(args):
     sys.stdout.write(protocol.armor(result, 'PUBLIC KEY BLOCK'))
 
 
-def main_create():
-    """Main function for GPG identity creation."""
-    p = argparse.ArgumentParser()
-    p.add_argument('user_id')
-    p.add_argument('-e', '--ecdsa-curve', default='nist256p1')
-    p.add_argument('-t', '--time', type=int, default=int(time.time()))
-    p.add_argument('-v', '--verbose', default=0, action='count')
-    p.add_argument('-s', '--subkey', default=False, action='store_true')
-
-    args = p.parse_args()
+def run_create(args):
+    """Export public GPG key."""
     util.setup_logging(verbosity=args.verbose)
     log.warning('This GPG tool is still in EXPERIMENTAL mode, '
                 'so please note that the API and features may '
@@ -82,13 +74,20 @@ def main_create():
     existing_gpg = keyring.gpg_version().decode('ascii')
     required_gpg = '>=2.1.11'
     if semver.match(existing_gpg, required_gpg):
-        run_create(args)
+        export_public_key(args)
     else:
         log.error('Existing gpg2 has version "%s" (%s required)',
                   existing_gpg, required_gpg)
 
 
-def main_agent():
+def run_unlock(args):
+    """Unlock hardware device (for future interaction)."""
+    util.setup_logging(verbosity=args.verbose)
+    d = device.detect()
+    log.info('unlocked %s device', d)
+
+
+def run_agent(_):
     """Run a simple GPG-agent server."""
     home_dir = os.environ.get('GNUPGHOME', os.path.expanduser('~/.gnupg/trezor'))
     config_file = os.path.join(home_dir, 'gpg-agent.conf')
@@ -115,12 +114,29 @@ def main_agent():
                     log.exception('gpg-agent failed: %s', e)
 
 
-def auto_unlock():
-    """Automatically unlock first found device (used for `gpg-shell`)."""
-    p = argparse.ArgumentParser()
+def main():
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser()
+    subparsers = parser.add_subparsers()
+
+    p = subparsers.add_parser('agent', help='Run GPG agent using a hardware device')
+    p.set_defaults(func=run_agent)
+
+    p = subparsers.add_parser('create', help='Export public GPG key')
+    p.add_argument('user_id')
+    p.add_argument('-e', '--ecdsa-curve', default='nist256p1')
+    p.add_argument('-t', '--time', type=int, default=int(time.time()))
     p.add_argument('-v', '--verbose', default=0, action='count')
+    p.add_argument('-s', '--subkey', default=False, action='store_true')
+    p.set_defaults(func=run_create)
 
-    args = p.parse_args()
-    util.setup_logging(verbosity=args.verbose)
-    d = device.detect()
-    log.info('unlocked %s device', d)
+    p = subparsers.add_parser('unlock', help='Unlock the hardware device')
+    p.add_argument('-v', '--verbose', default=0, action='count')
+    p.set_defaults(func=run_unlock)
+
+    args = parser.parse_args()
+    return args.func(args)
+
+
+if __name__ == '__main__':
+    main()