Replace magic number with magic dict

1 files changed, 9 insertions, 8 deletions

diff --git a/margfools b/margfools
index 7f6adfa..9524abe 100755
--- a/margfools
+++ b/margfools
@@ -23,16 +23,17 @@ def sign(data, key, pin=None, engine=None):
         env = None
     elif engine == 'pkcs11':
         # key on smartcard
+        digest_info = { # from RFC 3447
+            'MD2': '3020300c06082a864886f70d020205000410',
+            'MD5': '3020300c06082a864886f70d020505000410',
+            'SHA-1': '3021300906052b0e03021a05000414',
+            'SHA-256': '3031300d060960864801650304020105000420',
+            'SHA-384': '3041300d060960864801650304020205000430',
+            'SHA-512': '3051300d060960864801650304020305000440'
+        }
         cmd = ['pkcs11-tool', '--id', key, '-s', '-m', 'RSA-PKCS', '-p', 'env:PIN']
         env = {'PIN': pin}
-        """magic_prefix is ASN.1 DER for
-           DigestInfo ::= SEQUENCE {
-               digestAlgorithm DigestAlgorithm,
-               digest OCTET STRING
-           }
-        """
-        magic_prefix = bytes.fromhex("3031300d060960864801650304020105000420")
-        raw_data = magic_prefix + base64.b64decode(data)
+        raw_data = bytes.fromhex(digest_info['SHA-256']) + base64.b64decode(data)
     p = subprocess.run(cmd, env=env, input=raw_data, capture_output=True)
     return base64.b64encode(p.stdout).decode()