update dns-bl builder to use category-based responses

Author: superstes

src/riskdb/dnsbl/build_config.py

@@ -2,15 +2,55 @@
 
 from sys import argv
 from pathlib import Path
+from yaml import dump as yaml_dump
 from json import loads as json_loads
 
 # see: https://github.com/O-X-L/dnsbl-server
 
-# todo: separate response-codes for different report-categories
 FILE_DNSBL = '/tmp/riskdb-dnsbl.yml'
 
 INCLUDE_NET_REPUTATION = ['bad', 'warn']
 INCLUDE_IP_REPORTS = 1000
+INCLUDE_CAT_PERCENT = 0.33
+
+# risk-db cat to dns-bl cat
+CATEGORY_MAPPING = {
+    'bot': 'bot',
+    'probe': 'scanner',
+    'rate': 'scanner',
+    'attack': 'attack',
+    'crawler': 'bot',
+    'spam': 'spam',
+    'malware': 'attack',
+}
+# dns-bl cat to responses
+CATEGORY_TO_RESPONSE = {
+    'abuse': '127.0.0.2',
+    'scanner': '127.0.0.3',
+    'bot': '127.0.0.4',
+    'attack': '127.0.0.5',
+    'spam': '127.0.0.6',
+}
+
+
+def _get_categories(reports: dict) -> list[str]:
+    cats = []
+    s = reports['sum']
+    for c, v in reports.items():
+        if c == 'sum':
+            continue
+
+        if v > INCLUDE_IP_REPORTS:
+            cats.append(c)
+
+        elif v / s > INCLUDE_CAT_PERCENT:
+            cats.append(c)
+
+    cats = list(set(CATEGORY_MAPPING[c] for c in cats))
+    if len(cats) == 0:
+        cats = [CATEGORY_TO_RESPONSE['abuse']]
+
+    return cats
 
 
 def _to_yaml_list(d: list) -> str:
@@ -28,39 +68,49 @@ def main():
             data[k] = json_loads(f.read())
 
     print('BUILDING DNS-BL..')
-    dns_bl = {'nets': [], 'ips': []}
+    dns_bl = {
+        'abuse': {'nets': [], 'ips': []},
+        'scanner': {'nets': [], 'ips': []},
+        'bot': {'nets': [], 'ips': []},
+        'attack': {'nets': [], 'ips': []},
+        'spam': {'nets': [], 'ips': []},
+    }
 
     for ipp in ['net4', 'net6']:
         for net, net_info in data[ipp].items():
             if net_info['reputation'] in INCLUDE_NET_REPUTATION:
-                dns_bl['nets'].append(net)
+                for c in _get_categories(net_info['reports']):
+                    dns_bl[c]['nets'].append(net)
 
     data.pop('net4')
     data.pop('net6')
 
     for ipp in ['ip4', 'ip6']:
         for ip, ip_info in data[ipp].items():
             if ip_info['reports']['sum'] > INCLUDE_IP_REPORTS:
-                dns_bl['ips'].append(ip)
+                for c in _get_categories(ip_info['reports']):
+                    dns_bl[c]['ips'].append(ip)
 
     del data
 
+    data = {
+        'nets': [],
+        'ips': [],
+    }
+
+    for category, entries in dns_bl.items():
+        data['nets'].append({
+            'response': CATEGORY_TO_RESPONSE[category],
+            'content': entries['nets'],
+        })
+        data['ips'].append({
+            'response': CATEGORY_TO_RESPONSE[category],
+            'content': entries['ips'],
+        })
+
     print('WRITING DNS-BL..')
     with open(FILE_DNSBL, 'w', encoding='utf-8') as f:
-        f.write(f"""
----
-
-nets:
-  - response: 127.0.0.2
-    content:
-{_to_yaml_list(dns_bl['nets'])}
-
-ips:
-  - response: 127.0.0.2
-    content:
-{_to_yaml_list(dns_bl['ips'])}
-
-""")
+        f.write(yaml_dump(data))
 
     print('DONE:', FILE_DNSBL)