#!/usr/bin/env python3
"""Agent for implementing and auditing DMARC, DKIM, and SPF email security."""

import json
import argparse
import re
from datetime import datetime

try:
    import dns.resolver
except ImportError:
    dns = None


def check_spf_record(domain):
    """Check SPF record for a domain."""
    try:
        answers = dns.resolver.resolve(domain, "TXT")
        for rdata in answers:
            txt = rdata.to_text().strip('"')
            if txt.startswith("v=spf1"):
                issues = []
                if "+all" in txt:
                    issues.append({"issue": "spf_plus_all", "severity": "CRITICAL"})
                if "~all" in txt:
                    issues.append({"issue": "spf_softfail", "severity": "MEDIUM"})
                lookups = txt.count("include:") + txt.count("a:") + txt.count("mx")
                if lookups > 10:
                    issues.append({"issue": "spf_too_many_lookups", "count": lookups, "severity": "HIGH"})
                return {"domain": domain, "record": txt, "valid": True, "issues": issues}
        return {"domain": domain, "record": None, "valid": False,
                "issues": [{"issue": "no_spf_record", "severity": "HIGH"}]}
    except Exception as e:
        return {"domain": domain, "error": str(e)}


def check_dmarc_record(domain):
    """Check DMARC record for a domain."""
    try:
        answers = dns.resolver.resolve(f"_dmarc.{domain}", "TXT")
        for rdata in answers:
            txt = rdata.to_text().strip('"')
            if txt.startswith("v=DMARC1"):
                issues = []
                policy_match = re.search(r"p=(\w+)", txt)
                policy = policy_match.group(1) if policy_match else "none"
                if policy == "none":
                    issues.append({"issue": "dmarc_policy_none", "severity": "HIGH"})
                pct_match = re.search(r"pct=(\d+)", txt)
                if pct_match and int(pct_match.group(1)) < 100:
                    issues.append({"issue": "dmarc_pct_not_100", "severity": "MEDIUM"})
                if "rua=" not in txt:
                    issues.append({"issue": "no_aggregate_reporting", "severity": "MEDIUM"})
                return {"domain": domain, "record": txt, "policy": policy, "issues": issues}
        return {"domain": domain, "record": None,
                "issues": [{"issue": "no_dmarc_record", "severity": "CRITICAL"}]}
    except Exception as e:
        return {"domain": domain, "error": str(e)}


def check_dkim_record(domain, selector="default"):
    """Check DKIM record for a domain and selector."""
    try:
        dkim_domain = f"{selector}._domainkey.{domain}"
        answers = dns.resolver.resolve(dkim_domain, "TXT")
        for rdata in answers:
            txt = rdata.to_text().strip('"')
            issues = []
            if "k=rsa" in txt:
                p_match = re.search(r"p=([A-Za-z0-9+/=]+)", txt)
                if p_match and len(p_match.group(1)) < 300:
                    issues.append({"issue": "weak_dkim_key", "severity": "HIGH"})
            return {"domain": domain, "selector": selector, "record": txt[:200], "issues": issues}
        return {"domain": domain, "selector": selector, "record": None}
    except Exception as e:
        return {"domain": domain, "selector": selector, "error": str(e)}


def audit_domains(domain_list):
    """Audit multiple domains for email security records."""
    results = []
    for domain in domain_list:
        result = {
            "domain": domain,
            "spf": check_spf_record(domain),
            "dmarc": check_dmarc_record(domain),
            "dkim": check_dkim_record(domain),
        }
        all_issues = (result["spf"].get("issues", []) + result["dmarc"].get("issues", [])
                     + result["dkim"].get("issues", []))
        critical = sum(1 for i in all_issues if i.get("severity") == "CRITICAL")
        result["risk_level"] = "CRITICAL" if critical > 0 else "HIGH" if len(all_issues) > 2 else "MEDIUM"
        results.append(result)
    return results


def generate_dns_records(domain, policy="reject"):
    """Generate recommended SPF, DMARC, and DKIM DNS records."""
    return {
        "spf": f'v=spf1 include:_spf.google.com include:spf.protection.outlook.com -all',
        "dmarc": f'v=DMARC1; p={policy}; pct=100; rua=mailto:dmarc-reports@{domain}; '
                 f'ruf=mailto:dmarc-forensics@{domain}; adkim=s; aspf=s',
        "dkim_selector": "google._domainkey",
        "notes": "DKIM key must be generated by your email provider",
    }


def main():
    parser = argparse.ArgumentParser(description="DMARC/DKIM/SPF Email Security Agent")
    parser.add_argument("--domains", nargs="+", help="Domains to audit")
    parser.add_argument("--generate", help="Domain to generate records for")
    parser.add_argument("--output", default="email_security_report.json")
    args = parser.parse_args()

    report = {"generated_at": datetime.utcnow().isoformat(), "findings": {}}

    if args.domains:
        results = audit_domains(args.domains)
        report["findings"]["domain_audit"] = results
        for r in results:
            print(f"[+] {r['domain']}: {r['risk_level']}")

    if args.generate:
        records = generate_dns_records(args.generate)
        report["findings"]["recommended_records"] = records
        print(f"[+] Generated records for {args.generate}")

    with open(args.output, "w") as f:
        json.dump(report, f, indent=2, default=str)
    print(f"[+] Report saved to {args.output}")


if __name__ == "__main__":
    main()