#!/usr/bin/env python

# Add, remove or query IP addresses in DNSBL zone.
#
# Requires Debian package: python-dnspython

import argparse
import sys
import dns.update
import dns.query
import dns.tsigkeyring
import dns.resolver

__author__ = "Niccolo Rigacci"
__copyright__ = "Copyright 2020 Niccolo Rigacci <niccolo@rigacci.org>"
__license__ = "GPLv3-or-later"
__email__ = "niccolo@rigacci.org"
__version__ = "0.1.0"

# Set your DNSBL zone, DNS server and key here:
DNSBL_ZONE = 'bl.rigacci.org'
DNS_HOSTNAME = 'ns1.rigacci.org'
RNDC_KEY = {'bl-rigacci-org_rndc-key.' : 'rg2aizg+T6XkKkmpI42K7g=='}

def main():
    parser = argparse.ArgumentParser(description='Add, remove or query IP addresses in DNSBL zone.')
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('-a', '--add',    metavar='ADDRESS', help='add address to the zone')
    group.add_argument('-r', '--remove', metavar='ADDRESS', help='remove address from the zone')
    group.add_argument('-q', '--query',  metavar='ADDRESS', help='query for address in the zone')
    args = parser.parse_args()

    if args.query != None:
        address = args.query
        rev_address = '.'.join(reversed(address.split(".")))
        ret = QueryDNS(rev_address)
        if ret != None:
            print("Address %s is listed: %s.%s => %s" % (address, rev_address, DNSBL_ZONE, ret))
        else:
            print("Address %s is not listed" % (address,))

    elif args.add != None:
        address = args.add
        rev_address = '.'.join(reversed(address.split(".")))
        ret = AddDNS(rev_address)
        print ret

    elif args.remove != None:
        address = args.remove
        rev_address = '.'.join(reversed(address.split(".")))
        ret = DeleteDNS(rev_address)
        print ret

def AddDNS(rev_address, value='127.0.0.1'):
    print('Adding record type "A" for %s.%s' % (rev_address, DNSBL_ZONE))
    keyring = dns.tsigkeyring.from_text(RNDC_KEY)
    update = dns.update.Update(DNSBL_ZONE, keyring = keyring, keyalgorithm = 'hmac-md5.sig-alg.reg.int')
    update.add(rev_address, 8600, 'A', value)
    response = dns.query.tcp(update, DNS_HOSTNAME)
    return response

def DeleteDNS(rev_address):
    print('Removing record type "A" for %s.%s' % (rev_address, DNSBL_ZONE))
    keyring = dns.tsigkeyring.from_text(RNDC_KEY)
    update = dns.update.Update(DNSBL_ZONE, keyring = keyring, keyalgorithm = 'hmac-md5.sig-alg.reg.int')
    update.delete(rev_address, 'A')
    response = dns.query.tcp(update, DNS_HOSTNAME)
    return response

def QueryDNS(rev_address):
    query = rev_address + "." + DNSBL_ZONE
    resolver = dns.resolver.Resolver()
    resolver.timeout = 8
    resolver.lifetime = 8
    try:
        answers = resolver.query(query, "A")
        ret_val = answers[0]
    except:
        ret_val = None
    return ret_val

#---------------------------------------------------------------
# Main
#---------------------------------------------------------------
if __name__ == '__main__':
    main()