Fixup script structure for entrypoints

2018-03-26 03:35:03 -04:00
parent 30062b3c2b
commit e50d3907ae
6 changed files with 124 additions and 108 deletions
@@ -29,53 +29,57 @@ import argparse

 import dkim

-# Backward compatibility hack because argparse doesn't support optional
-# positional arguments
-arguments=['--'+arg if arg[:8] == 'identity' else arg for arg in sys.argv[1:]]
-parser = argparse.ArgumentParser(
-    description='Produce DKIM signature for email messages.',
-    epilog="message to be signed follows commands on stdin")
-parser.add_argument('selector', action="store")
-parser.add_argument('domain', action="store")
-parser.add_argument('privatekeyfile', action="store")
-parser.add_argument('--hcanon', choices=['simple', 'relaxed'],
-    default='relaxed',
-    help='Header canonicalization algorithm: default=relaxed')
-parser.add_argument('--bcanon', choices=['simple', 'relaxed'],
-    default='simple',
-    help='Body canonicalization algorithm: default=simple')
-parser.add_argument('--signalg', choices=['rsa-sha256', 'ed25519-sha256', 'rsa-sha1'],
-    default='rsa-sha256',
-    help='Signature algorithm: default=rsa-sha256')
-parser.add_argument('--identity', help='Optional value for i= tag.')
-args=parser.parse_args(arguments)
-include_headers = None
-length = None
-logger = None
+def main():
+    # Backward compatibility hack because argparse doesn't support optional
+    # positional arguments
+    arguments=['--'+arg if arg[:8] == 'identity' else arg for arg in sys.argv[1:]]
+    parser = argparse.ArgumentParser(
+        description='Produce DKIM signature for email messages.',
+        epilog="message to be signed follows commands on stdin")
+    parser.add_argument('selector', action="store")
+    parser.add_argument('domain', action="store")
+    parser.add_argument('privatekeyfile', action="store")
+    parser.add_argument('--hcanon', choices=['simple', 'relaxed'],
+        default='relaxed',
+        help='Header canonicalization algorithm: default=relaxed')
+    parser.add_argument('--bcanon', choices=['simple', 'relaxed'],
+        default='simple',
+        help='Body canonicalization algorithm: default=simple')
+    parser.add_argument('--signalg', choices=['rsa-sha256', 'ed25519-sha256', 'rsa-sha1'],
+        default='rsa-sha256',
+        help='Signature algorithm: default=rsa-sha256')
+    parser.add_argument('--identity', help='Optional value for i= tag.')
+    args=parser.parse_args(arguments)
+    include_headers = None
+    length = None
+    logger = None

-if sys.version_info[0] >= 3:
-    args.selector = bytes(args.selector, encoding='UTF-8')
-    args.domain = bytes(args.domain, encoding='UTF-8')
-    if args.identity is not None:
-        args.identity = bytes(args.identity, encoding='UTF-8')
-    args.hcanon = bytes(args.hcanon, encoding='UTF-8')
-    args.bcanon = bytes(args.bcanon, encoding='UTF-8')
-    args.signalg = bytes(args.signalg, encoding='UTF-8')
-    # Make sys.stdin and stdout binary streams.
-    sys.stdin = sys.stdin.detach()
-    sys.stdout = sys.stdout.detach()
-canonicalize = (args.hcanon, args.bcanon)
+    if sys.version_info[0] >= 3:
+        args.selector = bytes(args.selector, encoding='UTF-8')
+        args.domain = bytes(args.domain, encoding='UTF-8')
+        if args.identity is not None:
+            args.identity = bytes(args.identity, encoding='UTF-8')
+        args.hcanon = bytes(args.hcanon, encoding='UTF-8')
+        args.bcanon = bytes(args.bcanon, encoding='UTF-8')
+        args.signalg = bytes(args.signalg, encoding='UTF-8')
+        # Make sys.stdin and stdout binary streams.
+        sys.stdin = sys.stdin.detach()
+        sys.stdout = sys.stdout.detach()
+    canonicalize = (args.hcanon, args.bcanon)

-message = sys.stdin.read()
-try:
-    d = dkim.DKIM(message,logger=logger,
-                  signature_algorithm=args.signalg)
-    sig = d.sign(args.selector, args.domain, open(
-                 args.privatekeyfile, "rb").read(), identity = args.identity,
-                 canonicalize=canonicalize, include_headers=include_headers,
-                 length=length)
-    sys.stdout.write(sig)
-    sys.stdout.write(message)
-except Exception as e:
-    print(e, file=sys.stderr)
-    sys.stdout.write(message)
+    message = sys.stdin.read()
+    try:
+        d = dkim.DKIM(message,logger=logger,
+                      signature_algorithm=args.signalg)
+        sig = d.sign(args.selector, args.domain, open(
+                     args.privatekeyfile, "rb").read(), identity = args.identity,
+                     canonicalize=canonicalize, include_headers=include_headers,
+                     length=length)
+        sys.stdout.write(sig)
+        sys.stdout.write(message)
+    except Exception as e:
+        print(e, file=sys.stderr)
+        sys.stdout.write(message)
+
+if __name__ == "__main__":
+    main()