eulaurarien/feed_rules.py

53 lines
1.3 KiB
Python
Raw Normal View History

#!/usr/bin/env python3
import database
import argparse
import sys
import time
FUNCTION_MAP = {
'zone': database.Database.set_zone,
'ip4network': database.Database.set_ip4network,
'asn': database.Database.set_asn,
}
if __name__ == '__main__':
# Parsing arguments
parser = argparse.ArgumentParser(
description="TODO")
parser.add_argument(
'type',
choices=FUNCTION_MAP.keys(),
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()
fun = FUNCTION_MAP[args.type]
source: database.RulePath
if args.first_party:
source = database.RuleFirstPath()
else:
source = database.RuleMultiPath()
for rule in args.input:
rule = rule.strip()
try:
fun(DB,
rule,
source=source,
updated=int(time.time()),
)
except ValueError:
DB.log.error(f"Could not add rule: {rule}")
DB.save()