Simple script to strip all headers which are not included in the signature of a DKIM signed email

dkim-email-repair.py

import re
import copy
import logging
import argparse
import email.parser

logger = logging.getLogger()
logger.setLevel(logging.DEBUG)

required_headers = ["dkim-signature", "content-type"]
auth_header_keywords = ["spf", "authentication", "dkim"]


class DKIMSignature(object):
    def __init__(self, header_value):

        self.raw_header = header_value
        self.dkim_parts = self.parse_dkim_parameters()
        if self.dkim_parts.get("h"):
            # Split the list of signed headers from the DKIM signature
            self.signed_headers = [h.lower() for h in self.dkim_parts["h"].split(":")]
        else:
            self.signed_headers = None

    def parse_dkim_parameters(self):
        """
        Return dictionary of all options in the DKIM header.
        """
        params = re.split(r"\s*;\s*", self.raw_header)

        dkim_parts = {}
        for param in params:
            clean_param = re.sub(r"\n|\s", "", param, flags=re.MULTILINE)
            try:
                key, value = clean_param.split("=", 1)
                dkim_parts[key] = value
            except Exception:
                logging.debug("Invalid format for signature part: %s", clean_param)
        return dkim_parts


def email_with_only_signed_components(message, dkim_signature, include_debug=False):
    """
    Returns an email only including the header fields which are included in the signature.

    If include_debug is True then unsigned headers related to DKIM and SPF verification
    will also be included.
    """
    msg = copy.deepcopy(message)
    unsigned_fields = [f.lower() for f in msg.keys() if f.lower() not in dkim_signature.signed_headers]

    for field in unsigned_fields:
        # Include unsigned fields needed to create a valid DKIM signed email.
        if field not in required_headers:
            if include_debug and any(keyword in field for keyword in auth_header_keywords):
                continue
            logging.debug("Deleting unsigned field %s", field)
            del msg[field]
    return msg.as_string()


def parse_dkim(message, dkim_signature):
    """
    Perform rough validation of email message based on a DKIM signature
    """
    # Check that all signed fields are present in the email
    fields_missing = False
    for field in dkim_signature.signed_headers:
        if not field in message:
            logging.error("Email missing signed header: %s", field)
            fields_missing = True
    return not fields_missing


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Fix DKIM signatures in emails.')
    parser.add_argument('email', type=argparse.FileType('r'), help='The email file to parse')
    args = parser.parse_args()

    email_parser = email.parser.Parser()
    message = email_parser.parse(args.email)

    for dkim_value in message.get_all("DKIM-Signature", []):
        dkim_signature = DKIMSignature(dkim_value)
        result = parse_dkim(message, dkim_signature)
        if result:
            print(email_with_only_signed_components(message, dkim_signature, include_debug=True))
            break

    else:
        logging.error("No DKIM signatures found.")