DNSBlock

agent/agent.py

@@ -0,0 +1,126 @@
+import os
+import time
+import json
+import hashlib
+import requests
+import logging
+import subprocess
+from datetime import datetime
+from dotenv import load_dotenv
+
+# Load Env
+load_dotenv()
+
+API_URL = os.getenv('API_URL', 'http://localhost:8000/api')
+SERIAL_KEY = os.getenv('SERIAL_KEY')
+RPZ_FILE = os.getenv('RPZ_FILE', '/etc/unbound/rpz.zone')
+UNBOUND_CMD = os.getenv('UNBOUND_CMD', 'unbound-control reload')
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.FileHandler("agent.log"),
+        logging.StreamHandler()
+    ]
+)
+
+class DNSBlockAgent:
+    def __init__(self):
+        self.token = None
+        self.last_checksum = None
+
+    def authenticate(self):
+        try:
+            logging.info("Authenticating...")
+            response = requests.post(f"{API_URL}/auth/login", json={'serial_key': SERIAL_KEY})
+            response.raise_for_status()
+            data = response.json()
+            self.token = data['token']
+            logging.info("Authenticated successfully.")
+        except Exception as e:
+            logging.error(f"Authentication failed: {e}")
+            self.token = None
+
+    def get_domains(self):
+        if not self.token:
+            self.authenticate()
+            if not self.token:
+                return None
+
+        try:
+            headers = {'Authorization': f'Bearer {self.token}'}
+            response = requests.get(f"{API_URL}/v1/domains", headers=headers)
+            
+            if response.status_code == 401:
+                logging.warning("Token expired, re-authenticating...")
+                self.authenticate()
+                if self.token:
+                    headers = {'Authorization': f'Bearer {self.token}'}
+                    response = requests.get(f"{API_URL}/v1/domains", headers=headers)
+                else:
+                    return None
+
+            response.raise_for_status()
+            return response.json()
+        except Exception as e:
+            logging.error(f"Failed to fetch domains: {e}")
+            return None
+
+    def update_rpz(self, domains):
+        try:
+            logging.info(f"Updating RPZ file with {len(domains)} domains...")
+            
+            header = """$TTL 30m
+@ IN SOA localhost. root.localhost. (
+    {} ; Serial
+    3600       ; Refresh
+    1800       ; Retry
+    604800     ; Expire
+    86400 )    ; Minimum TTL
+    NS localhost.
+
+; Blocked Domains
+""".format(int(datetime.now().timestamp()))
+
+            with open(RPZ_FILE, 'w') as f:
+                f.write(header)
+                for domain in domains:
+                    f.write(f"{domain} CNAME .\n")
+                    f.write(f"*.{domain} CNAME .\n")
+            
+            logging.info("RPZ file updated.")
+            self.reload_unbound()
+        except Exception as e:
+            logging.error(f"Failed to update RPZ: {e}")
+
+    def reload_unbound(self):
+        try:
+            logging.info("Reloading Unbound...")
+            subprocess.run(UNBOUND_CMD.split(), check=True)
+            logging.info("Unbound reloaded.")
+        except Exception as e:
+            logging.error(f"Failed to reload Unbound: {e}")
+
+    def run(self):
+        logging.info("Agent started.")
+        while True:
+            data = self.get_domains()
+            if data:
+                checksum = data.get('checksum')
+                if checksum != self.last_checksum:
+                    logging.info(f"Checksum changed: {checksum}. Updating...")
+                    self.update_rpz(data.get('domains', []))
+                    self.last_checksum = checksum
+                else:
+                    logging.info("No changes detected.")
+            
+            time.sleep(300)
+
+if __name__ == "__main__":
+    if not SERIAL_KEY:
+        logging.error("SERIAL_KEY not defined in .env")
+        exit(1)
+    
+    agent = DNSBlockAgent()
+    agent.run()