#!/usr/bin/env python3
"""Threat intelligence collection agent using MISP/PyMISP.

Connects to MISP instances to collect, filter, and export threat intelligence
including events, attributes, and feeds via the PyMISP REST API client.
"""

import json
import os
import datetime

try:
    from pymisp import PyMISP
    HAS_PYMISP = True
except ImportError:
    HAS_PYMISP = False


def init_misp(url=None, key=None):
    """Initialize PyMISP client."""
    url = url or os.environ.get("MISP_URL", "https://misp.example.org")
    key = key or os.environ.get("MISP_API_KEY", "")
    if not HAS_PYMISP:
        return None
    return PyMISP(url, key, ssl=True)


def search_events(misp, tags=None, date_from=None, published=True, limit=50):
    """Search MISP events by tags and date."""
    if not misp:
        return {"error": "PyMISP not available"}
    kwargs = {"limit": limit, "published": published, "pythonify": True}
    if tags:
        kwargs["tags"] = tags
    if date_from:
        kwargs["date_from"] = date_from
    try:
        events = misp.search("events", **kwargs)
        return [
            {
                "id": e.id,
                "uuid": e.uuid,
                "info": e.info,
                "date": str(e.date),
                "threat_level": {1: "High", 2: "Medium", 3: "Low", 4: "Undefined"}.get(e.threat_level_id, "?"),
                "analysis": {0: "Initial", 1: "Ongoing", 2: "Complete"}.get(e.analysis, "?"),
                "attribute_count": e.attribute_count,
                "org": e.Orgc.name if hasattr(e, "Orgc") and e.Orgc else "",
                "tags": [t.name for t in (e.tags or [])],
            }
            for e in events
        ]
    except Exception as e:
        return {"error": str(e)}


def extract_attributes(misp, event_id, attr_type=None):
    """Extract attributes from a MISP event."""
    if not misp:
        return {"error": "PyMISP not available"}
    try:
        kwargs = {"eventid": event_id, "pythonify": True}
        if attr_type:
            kwargs["type_attribute"] = attr_type
        attrs = misp.search("attributes", **kwargs)
        return [
            {
                "type": a.type,
                "value": a.value,
                "category": a.category,
                "to_ids": a.to_ids,
                "comment": a.comment or "",
                "timestamp": str(datetime.datetime.fromtimestamp(int(a.timestamp))),
            }
            for a in attrs
        ]
    except Exception as e:
        return {"error": str(e)}


def collect_iocs_by_type(misp, ioc_types, date_from=None, limit=500):
    """Collect IOCs filtered by attribute type."""
    if not misp:
        return {"error": "PyMISP not available"}
    results = {}
    for ioc_type in ioc_types:
        try:
            kwargs = {"type_attribute": ioc_type, "to_ids": True, "pythonify": True, "limit": limit}
            if date_from:
                kwargs["date_from"] = date_from
            attrs = misp.search("attributes", **kwargs)
            results[ioc_type] = [
                {"value": a.value, "event_id": a.event_id, "comment": a.comment or ""}
                for a in attrs
            ]
        except Exception as e:
            results[ioc_type] = {"error": str(e)}
    return results


def list_feeds(misp):
    """List configured MISP feeds."""
    if not misp:
        return {"error": "PyMISP not available"}
    try:
        feeds = misp.feeds()
        return [
            {
                "id": f["Feed"]["id"],
                "name": f["Feed"]["name"],
                "provider": f["Feed"]["provider"],
                "url": f["Feed"]["url"],
                "enabled": f["Feed"]["enabled"],
                "source_format": f["Feed"]["source_format"],
            }
            for f in feeds
        ]
    except Exception as e:
        return {"error": str(e)}


def export_stix2(misp, event_id):
    """Export MISP event as STIX 2.1 bundle."""
    if not misp:
        return {"error": "PyMISP not available"}
    try:
        stix_data = misp.get_stix_event(event_id)
        return stix_data
    except Exception as e:
        return {"error": str(e)}


COMMON_IOC_TYPES = [
    "ip-dst", "ip-src", "domain", "hostname", "url",
    "md5", "sha1", "sha256", "email-src", "filename",
]


if __name__ == "__main__":
    print("=" * 60)
    print("Threat Intelligence Collection with MISP")
    print("PyMISP REST client, event search, attribute extraction, feeds")
    print("=" * 60)
    print("  PyMISP available: {}".format(HAS_PYMISP))

    misp = init_misp() if HAS_PYMISP else None

    if not misp:
        print("\n[DEMO] No MISP connection. Showing IOC types and feed structure.")
        print("\n--- Common IOC Types ---")
        for t in COMMON_IOC_TYPES:
            print("  - {}".format(t))
        print("\n--- Usage ---")
        print("  Set MISP_URL and MISP_API_KEY environment variables")
        print("  python agent.py")
    else:
        print("\n[*] Searching recent events...")
        events = search_events(misp, date_from="7d")
        if isinstance(events, list):
            print("  Found {} events".format(len(events)))
            for e in events[:5]:
                print("    [{}] {} ({} attrs)".format(e["id"], e["info"][:60], e["attribute_count"]))
        else:
            print("  Error: {}".format(events))

        feeds = list_feeds(misp)
        if isinstance(feeds, list):
            print("\n--- Feeds ({}) ---".format(len(feeds)))
            for f in feeds[:10]:
                status = "enabled" if f["enabled"] else "disabled"
                print("  [{}] {} ({})".format(f["id"], f["name"], status))

    print("\n" + json.dumps({"pymisp_available": HAS_PYMISP}, indent=2))