import json
import ssl
import socket
import datetime
import urllib.request
import http.client
import re
from concurrent.futures import ThreadPoolExecutor, as_completed

try:
    import dns.resolver as dnsresolver
    DNS_OK = True
except ImportError:
    DNS_OK = False

DISTRICTS_FILE = "districts.json"
RESULTS_FILE   = "results.json"
TIMEOUT        = 5
MAX_WORKERS    = 6

UA = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"

WEIGHTS = {
    "https_accessible":   13,
    "cert_valid":         13,
    "https_enforced":      9,
    "hsts":                9,
    "tls_version":         4,
    "x_frame_options":     3,
    "x_content_type":      2,
    "csp_header":          4,
    "permissions_policy":  3,
    "referrer_policy":     3,
    "cookies_secure_httponly": 4,
    "robots_txt":          3,
    "no_server_banner":    4,
    "spf":                 9,
    "dmarc":               9,
    "dnssec":              4,
    "no_zone_transfer":    4,
}

MAX_SCORE = sum(WEIGHTS.values())

def check_https(domain):
    r = {
        "https_accessible": False,
        "cert_valid":        False,
        "cert_expiry":      None,
        "cert_days_left":    None,
        "tls_version":      False,
        "tls_version_str":  None,
        "hsts":              False,
        "hsts_value":        None,
        "no_server_banner": False,
        "server_header":    None,
        "x_frame_options":  False,
        "x_content_type":    False,
        "csp_header":        False,
        "permissions_policy": False,
        "referrer_policy":  False,
        "cookies_secure_httponly": False,
        "robots_txt":        False,
    }
    
    def check_cookies_secure_httponly(resp):
        cookies = resp.headers.get_all('Set-Cookie') if hasattr(resp.headers, 'get_all') \
                  else [resp.headers.get('Set-Cookie', '')]

        if not any(c.strip() for c in cookies):
            return False

        for cookie_str in cookies:
            if not cookie_str:
                continue
            parts = [p.strip().lower() for p in cookie_str.split(';')]
            has_secure = 'secure' in parts
            has_httponly = 'httponly' in parts
            if not (has_secure and has_httponly):
                return False
        return True
    
    try:
        ctx = ssl.create_default_context()
        with socket.create_connection((domain, 443), timeout=TIMEOUT) as sock:
            with ctx.wrap_socket(sock, server_hostname=domain) as s:
                r["https_accessible"] = True
                r["tls_version_str"]  = s.version()
                r["tls_version"]      = s.version() in ("TLSv1.2", "TLSv1.3")
                cert = s.getpeercert()
                r["cert_valid"] = True
                exp = cert.get("notAfter", "")
                if exp:
                    dt = datetime.datetime.strptime(exp, "%b %d %H:%M:%S %Y %Z")
                    r["cert_expiry"]    = dt.strftime("%Y-%m-%d")
                    now = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
                    r["cert_days_left"] = (dt - now).days

        req = urllib.request.Request("https://" + domain, headers={"User-Agent": UA}, method="HEAD")
        with urllib.request.urlopen(req, timeout=TIMEOUT) as resp:
            h = {k.lower(): v for k, v in resp.headers.items()}
            r["hsts"] = bool(h.get("strict-transport-security", ""))
            r["hsts_value"] = h.get("strict-transport-security") or None
            sv = h.get("server", "")
            r["server_header"] = sv
            r["no_server_banner"] = not any(c.isdigit() for c in sv)
            r["x_frame_options"] = "x-frame-options" in h
            r["x_content_type"] = h.get("x-content-type-options", "").lower() == "nosniff"
            r["csp_header"] = "content-security-policy" in h
            r["permissions_policy"] = bool(h.get("permissions-policy", "") or h.get("feature-policy", ""))
            r["referrer_policy"] = "referrer-policy" in h
            r["cookies_secure_httponly"] = check_cookies_secure_httponly(resp)
    except ssl.SSLCertVerificationError:
        r["https_accessible"] = True
        r["cert_valid"] = False
    except Exception:
        pass

    try:
        req2 = urllib.request.Request("https://" + domain + "/robots.txt", headers={"User-Agent": UA})
        with urllib.request.urlopen(req2, timeout=TIMEOUT) as resp:
            body = resp.read(2048).decode("utf-8", errors="ignore").lower()
            r["robots_txt"] = resp.status == 200 and "disallow" in body
    except Exception:
        pass
    return r

def check_redirect(domain):
    r = {"https_redirect": False, "http_closed": False, "https_enforced": False, "http_status": None, "redirect_location": None}
    try:
        c = http.client.HTTPConnection(domain, timeout=TIMEOUT)
        c.request("HEAD", "/", headers={"User-Agent": UA})
        resp = c.getresponse()
        loc = resp.getheader("Location", "") or ""
        r["http_status"], r["redirect_location"] = resp.status, loc
        r["https_redirect"] = resp.status in (301, 302, 307, 308) and loc.startswith("https://")
        c.close()
    except Exception:
        r["http_closed"] = True
    r["https_enforced"] = r["https_redirect"] or r["http_closed"]
    return r

def check_dns(domain):
    r = {"spf": False, "spf_value": None, "dmarc": False, "dmarc_value": None, "dmarc_policy": None, "dnssec": False, "no_zone_transfer": False, "mx_exists": False, "ip": None}
    try:
        answers = dnsresolver.Resolver()
        answers.nameservers = ['8.8.8.8', '1.1.1.1']
        r["ip"] = str(answers.resolve(domain, 'A')[0])
    except Exception:
        try:
            r["ip"] = socket.gethostbyname(domain)
        except Exception:
            pass
    if not DNS_OK: return r
    base_domain = domain[4:] if domain.startswith("www.") else domain
    res = dnsresolver.Resolver()
    res.nameservers = ['8.8.8.8', '1.1.1.1']  # ← Публичен DNS!
    res.lifetime = TIMEOUT
    try:
        for rd in res.resolve(base_domain, "TXT"):
            t = rd.to_text().strip('"')
            if t.startswith("v=spf1"): r["spf"], r["spf_value"] = True, t; break
    except Exception: pass
    try:
        for rd in res.resolve("_dmarc." + base_domain, "TXT"):
            t = rd.to_text().strip('"')
            if t.startswith("v=DMARC1"):
                r["dmarc"], r["dmarc_value"] = True, t
                for p in t.split(";"):
                    if p.strip().startswith("p="): r["dmarc_policy"] = p.strip()[2:]
                break
    except Exception: pass
    try:
        res.resolve(base_domain, "MX")
        r["mx_exists"] = True
    except Exception: pass
    try:
        ans = res.resolve(base_domain, "RRSIG", raise_on_no_answer=False)
        if ans.rrset: r["dnssec"] = True
    except Exception: pass
    try:
        res.resolve(base_domain, "AXFR")  # Ако успее → уязвим
        r["no_zone_transfer"] = False     # Zone transfer разрешен → ГРЕШНО
    except Exception:
        r["no_zone_transfer"] = True      # Zone transfer блокиран → ПРАВИЛНО
    
    return r

def calc_score(checks):
    return sum(WEIGHTS[k] for k in WEIGHTS if checks.get(k))

def get_color(score):
    if score >= 90: return "#2ecc71"
    if score >= 60: return "#f1c40f"
    if score >= 35: return "#e67e22"
    return "#e74c3c"

def get_status(score):
    if score >= 90: return "good"
    if score >= 60: return "warning"
    if score >= 35: return "poor"
    return "critical"

def scan_domain(d):
    dom = d["domain"]
    print("  " + dom)
    checks = {}
    checks.update(check_dns(dom))
    checks.update(check_https(dom))
    checks.update(check_redirect(dom))
    score = calc_score(checks)
    now = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d %H:%M UTC")
    return {
        "id": d["id"], "name": d["name"], "domain": dom, "lat": d["lat"], "lon": d["lon"],
        "special": d.get("special", False), "score": score, "max_score": MAX_SCORE,
        "color": get_color(score), "status": get_status(score), "scanned_at": now, "checks": checks,
    }

# ── MAIN ──────────────────────────────────────────────────────────
if __name__ == "__main__":
    try:
        with open(DISTRICTS_FILE, encoding="utf-8") as f:
            districts = json.load(f)
        
        print(f"\nСканиране на {len(districts)} домейна (max: {MAX_SCORE} т.)...\n")
        
        results = []
        with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
            futures = {executor.submit(scan_domain, d): d for d in districts}
            for future in as_completed(futures):
                try:
                    results.append(future.result())
                except Exception as e:
                    print(f"  Грешка при {futures[future]['domain']}: {e}")

        results.sort(key=lambda x: x["id"])
        with open(RESULTS_FILE, "w", encoding="utf-8") as f:
            json.dump(results, f, ensure_ascii=False, indent=2)

        print("\n" + "Район".ljust(20) + " " + "Домейн".ljust(32) + " Score  Статус")
        print("-" * 70)
        icons = {"good": "[+]", "warning": "[~]", "poor": "[!]", "critical": "[X]"}
        for r in results:
            icon = icons.get(r["status"], "[ ]")
            print(f"{r['name'].ljust(20)} {r['domain'].ljust(32)} {str(r['score']).rjust(5)}  {icon} {r['status']}")

        # ── Вграждане в HTML ──────────────────────────────────────────────
        HTML_FILE = "sofia_security_map.html"
        try:
            with open(HTML_FILE, "r", encoding="utf-8") as f:
                html = f.read()
            results_json = json.dumps(results, ensure_ascii=False)
            marker_start, marker_end = "// DATA_START", "// DATA_END"
            new_block = f"{marker_start}\nconst SCAN_DATA = {results_json};\n{marker_end}"
            if marker_start in html:
                html = re.sub(r'// DATA_START.*?// DATA_END', new_block, html, flags=re.DOTALL)
            else:
                html = html.replace("// DATA_INJECT_HERE", new_block)
            with open(HTML_FILE, "w", encoding="utf-8") as f:
                f.write(html)
            print("\nКартата е обновена: " + HTML_FILE)
        except Exception as e:
            print("Грешка при обновяване на HTML: " + str(e))

    except Exception as e:
        print(f"Критична грешка: {e}")

    input("\nНатисни Enter за изход...")