#!/usr/bin/env python3
#
# 
# This script takes two arguments: 
#  1) a BIND DNS zone file
#  2) a DNS record generated by 'opendkim-genkey'
#
# 'opendkim-genkey' is part of the 'opendkim-tools' Debian package
#
# This script use regex to 
#  1) extract the DKIM selector and the value of the TXT record from the
#     DKIM TXT record generated by opendkim-genkey
#  2) replace the DKIM selector and TXT value in the zone file with the
#     new values
#
#  Example:
#  First run:
#   /usr/sbin/opendkim-genkey -b 2048 -d example.org -s $(/bin/date +%Y%m%d)-1
#  which will generate a DNS DKIM TXT record such as '20210811-1.txt'
#  Then run this script, passing the zone file you want to update and the above
#    .txt file:
#   update-dkim.py <zone file> 20210811-1.txt
#

import sys

# Make sure Python is at least version 3.6 (required for f-string literals support)
if not (sys.version_info.major == 3 and sys.version_info.minor >= 6):
    print("This script requires Python 3.6 or higher")
    print("You are using Python {}.{}.".format(sys.version_info.major, sys.version_info.minor))
    sys.exit(1)

import argparse
import re

# Create an ArgumentParser object to parse the command-line arguments
parser = argparse.ArgumentParser(
    description='Update DNS zone file DKIM TXT record(s)'
)

# First argument is the zone file which will be updated    
parser.add_argument(
    'zone_file', type=argparse.FileType('r+')
)
    
# Second argument is the opendkim-genkey-generated DKIM TXT record
parser.add_argument(
    'dkim_txt', type=argparse.FileType('r')
)

args = parser.parse_args()

# Read the files into corresponding string variables
zone_text = args.zone_file.read()
dkim_txt = args.dkim_txt.read()

# This regex strips the old DKIM selector and TXT value from the zone file
# The remaining text is captured in a backreference
zone_re = re.compile(
    r'^[\d-]+(\._domainkey\.?[a-z.]*\s+[0-9h]+\s+IN\s+TXT\s+\()[\s"=;+\/\w]+',
    re.MULTILINE
)

# This regex uses two backreferences to capture the DKIM selector and TXT value
#  from the opendkim-genkey-generated DKIM record
dkim_txt_re = re.compile(
    r'^([\d-]+)\._domainkey\s+IN\s+TXT\s+\(([\s"=;+\/\w]+).+\s+$',
    re.MULTILINE
)

# Store the captured values as variables
selector = dkim_txt_re.sub(r'\1', dkim_txt)
dkim_value = dkim_txt_re.sub(r'\2', dkim_txt)

# Perform the regex substitution on the zone file
#  use f-strings to provide the variables in the substitution
zone = zone_re.sub(rf"{selector}\1{dkim_value}", zone_text)

# print("\nselector: ", selector, "\ndkim_value: ", dkim_value, "\n\nzone:\n", zone)
print(zone)

# Sample DKIM TXT record
#
# 20160525-014646._domainkey.whitehall 3h IN TXT ( "v=DKIM1; k=rsa; t=y; p=MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCA"
#     "QEAulkdTaAsWGpcN4O6HRMzCN6i67AWoMyjemrryvd+j/2epCSZ3qNfkD/ZV8g4S"
#     "42IVk4dBs3evQQCyCGQhhoqpQIDAQAB" )