Calculate DS record from DNSKEY with Python 3

dnskey_to_dsrecord.py

#!/usr/bin/env python3

# vim: autoindent tabstop=4 shiftwidth=4 expandtab softtabstop=4 filetype=python

"""
Generate a DNSSEC DS record based on the existing DNSKEY record


usage: dnskey_to_dsrecord.py [-h] DOMAIN-NAME

DNSSEC DS-key Calculator

positional arguments:
  DOMAIN-NAME  Domain to query for

options:
  -h, --help   show this help message and exit


Author: Wido den Hollander <[email protected]>
Improvements: Jari Turkia <[email protected]>

Many thanks to this blogpost: https://www.v13.gr/blog/?p=239

Source: https://gist.github.com/wido/4c6288b2f5ba6d16fce37dca3fc2cb4a
"""

import struct
import base64
import hashlib
import DNS  # pip install py3dns
import argparse


def _calc_keyid(flags, protocol, algorithm, dnskey) -> int:
    st = struct.pack('!HBB', int(flags), int(protocol), int(algorithm))
    st += base64.b64decode(dnskey)

    cnt = 0
    for idx in range(len(st)):
        s = struct.unpack('B', st[idx:idx + 1])[0]
        if (idx % 2) == 0:
            cnt += s << 8
        else:
            cnt += s

    return ((cnt & 0xFFFF) + (cnt >> 16)) & 0xFFFF


def _calc_ds(domain, flags, protocol, algorithm, dnskey) -> dict[str, str]:
    if not domain.endswith('.'):
        domain += '.'

    signature = bytes()
    for i in domain.split('.'):
        signature += struct.pack('B', len(i)) + i.encode()

    signature += struct.pack('!HBB', int(flags), int(protocol), int(algorithm))
    signature += base64.b64decode(dnskey)

    return {
        'sha1': hashlib.sha1(signature).hexdigest().upper(),
        'sha256': hashlib.sha256(signature).hexdigest().upper(),
    }


def dnskey_to_ds(domain, flags, protocol, algorithm, key_base64) -> dict[str, str]:
    keyid = _calc_keyid(flags, protocol, algorithm, key_base64)
    ds = _calc_ds(domain, flags, protocol, algorithm, key_base64)

    ret = {
        'sha1': "{0:d} {1:d} {2:d} {3:s}".format(keyid, algorithm, 1, ds['sha1'].lower()),
        'sha256': "{0:d} {1:d} {2:d} {3:s}".format(keyid, algorithm, 2, ds['sha256'].lower())
    }

    return ret


def dns_request(domain: str) -> tuple[int, int, int, str]:
    opts = {
        'qtype': "DNSKEY"
    }
    reqobj = DNS.Request()
    res = reqobj.req(domain, **opts)
    if not res.answers:
        raise ValueError("No DNSKEY for '{}'!".format(domain))

    # Post-process received binary data
    data = res.answers[0]['data']
    encoded = base64.b64encode(data[4:])

    return 257, data[2], data[3], encoded.decode('ascii')


def main() -> None:
    parser = argparse.ArgumentParser(description='DNSSEC DS-key Calculator')
    parser.add_argument('domain',
                        metavar="DOMAIN-NAME",
                        help="Domain to query for")
    args = parser.parse_args()

    flags, protocol, algorithm, public_key = dns_request(args.domain)

    ds_data = dnskey_to_ds(args.domain, flags, protocol, algorithm, public_key)
    print("{0:s}, {1:d} {2:d} {3:d} DS-records:".format(args.domain, flags, protocol, algorithm))
    print("SHA-1 hashed    : {}".format(ds_data['sha1']))
    print("SHA-2 256 hashed: {}".format(ds_data['sha256']))


if __name__ == "__main__":
    main()