eulaurarien/feed_rules.py

#!/usr/bin/env python3

import database
import argparse
import sys
import ipaddress


if __name__ == '__main__':

    # Parsing arguments
    parser = argparse.ArgumentParser(
        description="TODO")
    parser.add_argument(
        'type',
        choices={'zone', 'ip4network'},
        help="Type of rule inputed")
    parser.add_argument(
        '-i', '--input', type=argparse.FileType('r'), default=sys.stdin,
        help="List of domains domains to block (with their subdomains)")
    parser.add_argument(
        '-f', '--first-party', action='store_true',
        help="The input only comes from verified first-party sources")
    args = parser.parse_args()

    DB = database.Database()

    FUNCTION_MAP = {
            'zone': DB.set_zone,
            'ip4network': DB.set_ip4network,
    }

    fun = FUNCTION_MAP[args.type]

    for rule in args.input:
        fun(rule.strip(), is_first_party=args.first_party)

    DB.close()