registry.py

import re
import yaml
import orjson
import requests
from pathlib import Path
from typing import Dict, List, Optional, Set
from loguru import logger
from rapidfuzz import fuzz
from dataclasses import dataclass
from config import DATA_DIR, COUNTRIES

logger.add(DATA_DIR / "app.log", rotation="1 MB", retention="7 days")

@dataclass
class Intent:
    name: str
    keywords: List[str]
    defaults: Dict[str, str]
    platform_templates: Dict[str, List[str]]

@dataclass
class CandidateStore:
    intents: Dict[str, Intent]
    services: List[Dict[str, str]]
    countries: Set[str]
    operators: Set[str]

class Registry:
    def __init__(self):
        self.store = CandidateStore(
            intents={},
            services=[],
            countries=set(COUNTRIES),
            operators={"AND", "OR", "NOT", "after", "before", "city", "country", "geo", "hostname", "ip", "net", "org", "os", "port", "product", "title"}
        )
        self.load_intents()
        self.load_services()
        self.load_ingested()

    def load_intents(self):
        intents_file = DATA_DIR / "intents.yaml"
        if intents_file.exists():
            with open(intents_file, "r") as f:
                data = yaml.safe_load(f)
                for name, details in data.get("intents", {}).items():
                    self.store.intents[name] = Intent(
                        name=name,
                        keywords=details.get("keywords", []),
                        defaults=details.get("defaults", {}),
                        platform_templates=details.get("platform_templates", {})
                    )
            logger.info(f"Loaded {len(self.store.intents)} intents from {intents_file}")

    def load_services(self):
        services_file = DATA_DIR / "services.yaml"
        if services_file.exists():
            with open(services_file, "r") as f:
                self.store.services = yaml.safe_load(f)
            logger.info(f"Loaded {len(self.store.services)} services from {services_file}")

    def load_ingested(self):
        ingested_dir = DATA_DIR / "ingested"
        ingested_dir.mkdir(exist_ok=True)
        for file in ingested_dir.glob("*.json"):
            with open(file, "rb") as f:
                data = orjson.loads(f.read())
                for item in data.get("intents", []):
                    self.store.intents[item["name"]] = Intent(**item)
                for item in data.get("services", []):
                    if item not in self.store.services:
                        self.store.services.append(item)
            logger.info(f"Ingested data from {file}")

    def ingest_from_url(self, url: str):
        try:
            response = requests.get(url, timeout=10)
            response.raise_for_status()
            data = response.json()
            ingested_dir = DATA_DIR / "ingested"
            ingested_dir.mkdir(exist_ok=True)
            file_path = ingested_dir / f"ingested_{url.split('/')[-1]}.json"
            with open(file_path, "wb") as f:
                f.write(orjson.dumps(data))
            self.load_ingested()
            logger.info(f"Ingested data from {url}")
        except Exception as e:
            logger.error(f"Failed to ingest from {url}: {e}")

    def parse_query(self, query: str) -> Dict[str, any]:
        tokens = re.findall(r'"[^"]*"|[^\s"]+', query)
        parsed = {
            "raw": query,
            "intents": [],
            "services": [],
            "filters": {},
            "operators": []
        }
        current_operator = None
        for token in tokens:
            if token in self.store.operators:
                current_operator = token
                parsed["operators"].append(token)
            elif token.lower() in self.store.countries:
                parsed["filters"]["country"] = token.lower()
            elif re.match(r"^\d{1,5}$", token):
                parsed["filters"]["port"] = token
            elif re.match(r"^[a-zA-Z0-9-]+\.[a-zA-Z]{2,}$", token):
                parsed["filters"]["hostname"] = token
            elif re.match(r"^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}(?:/\d{1,2})?$", token):
                parsed["filters"]["ip" if "/" not in token else "net"] = token
            else:
                intent_match = max(
                    ((name, fuzz.partial_ratio(token.lower(), k)) for name, intent in self.store.intents.items() for k in intent.keywords),
                    key=lambda x: x[1],
                    default=(None, 0)
                )
                if intent_match[1] > 80:
                    parsed["intents"].append(intent_match[0])
                service_match = max(
                    ((s["name"], fuzz.partial_ratio(token.lower(), s["name"].lower() + " " + " ".join(s.get("aliases", [])))) for s in self.store.services),
                    key=lambda x: x[1],
                    default=(None, 0)
                )
                if service_match[1] > 80:
                    parsed["services"].append(service_match[0])
                elif current_operator in {"os", "org", "city", "title"}:
                    parsed["filters"][current_operator] = token.strip('"')
        return parsed

    def sanitize_operators(self, query: str) -> str:
        tokens = re.findall(r'"[^"]*"|[^\s"]+', query)
        sanitized = []
        for token in tokens:
            if token in self.store.operators or token.lower() in self.store.countries or re.match(r"^\d{1,5}$", token) or re.match(r"^[a-zA-Z0-9-]+\.[a-zA-Z]{2,}$", token) or re.match(r"^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}(?:/\d{1,2})?$", token):
                sanitized.append(token)
            else:
                sanitized.append(f'"{token}"')
        return " ".join(sanitized)