Signature algorithm names are now defined in dkim.crypto.

dkim/__init__.py

@@ -28,6 +28,7 @@ import time
 from dkim.canonicalization import algorithms
 from dkim.crypto import (
     DigestTooLargeError,
+    HASH_ALGORITHMS,
     parse_pem_private_key,
     parse_public_key,
     RSASSA_PKCS1_v1_5_sign,
@@ -202,8 +203,9 @@ def fold(header):
 
 
 def sign(message, selector, domain, privkey, identity=None,
-         canonicalize=(b'simple', b'simple'), include_headers=None, length=False,
+         canonicalize=(b'simple', b'simple'),
-         logger=None):
+         signature_algorithm=b'rsa-sha256',
+         include_headers=None, length=False, logger=None):
     """Sign an RFC822 message and return the DKIM-Signature header line.
 
     @param message: an RFC822 formatted message (with either \\n or \\r\\n line endings)
@@ -245,7 +247,7 @@ def sign(message, selector, domain, privkey, identity=None,
 
     sigfields = [x for x in [
         (b'v', b"1"),
-        (b'a', b"rsa-sha256"),
+        (b'a', signature_algorithm),
         (b'c', b"/".join(
             (algorithms[canonicalize[0]].name,
              algorithms[canonicalize[1]].name))),
@@ -335,12 +337,10 @@ def verify(message, logger=None, dnsfunc=get_txt):
     headers = header_algorithm.canonicalize_headers(headers)
     body = body_algorithm.canonicalize_body(body)
 
-    if sig[b'a'] == b"rsa-sha1":
+    try:
-        hasher = hashlib.sha1
+        hasher = HASH_ALGORITHMS[sig[b'a']]
-    elif sig[b'a'] == b"rsa-sha256":
+    except KeyError as e:
-        hasher = hashlib.sha256
+        logger.error("unknown signature algorithm: %s" % e.message)
-    else:
-        logger.error("unknown signature algorithm (%s)" % sig[b'a'])
         return False
 
     if b'l' in sig:

dkim/crypto.py

@@ -21,6 +21,7 @@
 
 __all__ = [
     'DigestTooLargeError',
+    'HASH_ALGORITHMS',
     'parse_pem_private_key',
     'parse_private_key',
     'parse_public_key',
@@ -30,6 +31,7 @@ __all__ = [
     ]
 
 import base64
+import hashlib
 import re
 
 from dkim.asn1 import (
@@ -76,6 +78,10 @@ ASN1_RSAPrivateKey = [
     ])
 ]
 
+HASH_ALGORITHMS = {
+    b'rsa-sha1': hashlib.sha1,
+    b'rsa-sha256': hashlib.sha256,
+    }
 
 # These values come from RFC 3447, section 9.2 Notes, page 43.
 HASH_ID_MAP = {