changeset 7:47b4c98e4d40

Allow algorithm selection for all algorithms in :mod:`hashlib`. Besides of the name "shasum" the MD5, BLAKE2b, BLAKE2s and SHA-3 variants are implemented.
author Franz Glasner <fzglas.hg@dom66.de>
date Fri, 04 Dec 2020 15:16:20 +0100
parents a21e83c855cc
children 048b97213a23
files shasum.py
diffstat 1 files changed, 51 insertions(+), 4 deletions(-) [+]
line wrap: on
line diff
--- a/shasum.py	Fri Dec 04 13:15:53 2020 +0100
+++ b/shasum.py	Fri Dec 04 15:16:20 2020 +0100
@@ -27,6 +27,9 @@
         description="Python implementation of shasum",
         fromfile_prefix_chars='@')
     aparser.add_argument(
+        "--algorithm", "-a", action="store", type=argv2algo,
+        help="1, 224, 256 (default), 384, 512, 3-224, 3-256, 3-384, 3-512, blake2b, blake2s, md5")
+    aparser.add_argument(
         "--binary", "-b", action="store_false", dest="text_mode", default=False,
         help="read in binary mode (default)")
     aparser.add_argument(
@@ -44,6 +47,9 @@
         print("ERROR: text mode not supported", file=sys.stderr)
         sys.exit(78)   # :manpage:`sysexits(3)` EX_CONFIG
 
+    if not opts.algorithm:
+        opts.algorithm = argv2algo("256")
+
     if opts.bsd:
         out = out_bsd
     else:
@@ -60,20 +66,61 @@
         else:
             source = sys.stdin.buffer
         out(sys.stdout,
-            compute_digest(hashlib.sha256, source),
+            compute_digest(opts.algorithm[0], source),
             None,
-            "SHA256",
+            opts.algorithm[1],
             True)
     else:
         for fn in opts.files:
             with open(fn, "rb") as source:
                 out(sys.stdout,
-                    compute_digest(hashlib.sha256, source),
+                    compute_digest(opts.algorithm[0], source),
                     fn,
-                    "SHA256",
+                    opts.algorithm[1],
                     True)
 
 
+def argv2algo(s):
+    """Convert a commane line algorithm specifier into a tuple with the
+    type/factory of the digest and the algorithms tag for output purposes
+
+    :param str s: the specifier from the commane line
+    :return: the internal digest specification
+    :rtype: a tuple (digest_type_or_factory, name_in_output)
+
+    String comparisons are done case-insensitively.
+
+    """
+    s = s.lower()
+    if s in ("1", "sha1"):
+        return (hashlib.sha1, "SHA1")
+    elif s in ("224", "sha224"):
+        return (hashlib.sha224, "SHA224")
+    elif s in ("256", "sha256"):
+        return (hashlib.sha256, "SHA256")
+    elif s in ("384", "sha384"):
+        return (hashlib.sha384, "SHA384")
+    elif s in ("512", "sha512"):
+        return (hashlib.sha512, "SHA512")
+    elif s in ("3-224", "sha3-224"):
+        return (hashlib.sha3_224, "SHA3-224")
+    elif s in ("3-256", "sha3-256"):
+        return (hashlib.sha3_256, "SHA3-256")
+    elif s in ("3-384", "sha3-384"):
+        return (hashlib.sha3_384, "SHA3-384")
+    elif s in ("3-512", "sha3-512"):
+        return (hashlib.sha3_512, "SHA3-512")
+    elif s in ("blake2b", "blake2b-512"):
+        return (hashlib.blake2b, "BLAKE2b")
+    elif s in ("blake2s", "blake2s-256"):
+        return (hashlib.blake2s, "BLAKE2s")
+    elif s == "md5":
+        return (hashlib.md5, "MD5")
+    else:
+        raise argparse.ArgumentTypeError(
+            "`{}' is not a recognized algorithm".format(s))
+
+
 def out_bsd(dest, digest, filename, digestname, binary):
     if filename is None:
         print(digest, file=dest)