diff --git a/dkim/__init__.py b/dkim/__init__.py index dee2b7a..6311ab1 100644 --- a/dkim/__init__.py +++ b/dkim/__init__.py @@ -25,7 +25,10 @@ import logging import re import time -from dkim.canonicalization import algorithms +from dkim.canonicalization import ( + algorithms, + CanonicalizationPolicy, + ) from dkim.crypto import ( DigestTooLargeError, HASH_ALGORITHMS, @@ -334,8 +337,9 @@ def verify(message, logger=None, dnsfunc=get_txt): except KeyError as e: logger.error("unknown canonicalization algorithm: %s" % e.message) return False - headers = header_algorithm.canonicalize_headers(headers) - body = body_algorithm.canonicalize_body(body) + canon_policy = CanonicalizationPolicy(header_algorithm, body_algorithm) + headers = canon_policy.canonicalize_headers(headers) + body = canon_policy.canonicalize_body(body) try: hasher = HASH_ALGORITHMS[sig[b'a']] diff --git a/dkim/canonicalization.py b/dkim/canonicalization.py index 9b191f2..7023dc8 100644 --- a/dkim/canonicalization.py +++ b/dkim/canonicalization.py @@ -23,6 +23,7 @@ import re __all__ = [ 'algorithms', + 'CanonicalizationPolicy', ] @@ -83,4 +84,17 @@ class Relaxed: compress_whitespace(strip_trailing_whitespace(body))) +class CanonicalizationPolicy: + + def __init__(self, header_algorithm, body_algorithm): + self.header_algorithm = header_algorithm + self.body_algorithm = body_algorithm + + def canonicalize_headers(self, headers): + return self.header_algorithm.canonicalize_headers(headers) + + def canonicalize_body(self, body): + return self.body_algorithm.canonicalize_body(body) + + algorithms = dict((c.name, c) for c in (Simple, Relaxed))