#!/usr/bin/env python3
"""STIX/TAXII threat intelligence feed integration agent.

Connects to TAXII 2.0/2.1 servers to discover and consume threat intelligence
feeds in STIX format. Extracts indicators of compromise (IOCs), threat actors,
malware families, and attack patterns from STIX bundles.
"""
import argparse
import json
import os
import sys
from datetime import datetime, timezone

try:
    from taxii2client.v20 import Server as Server20, Collection as Collection20
    from taxii2client.v21 import Server as Server21, Collection as Collection21
    HAS_TAXII_CLIENT = True
except ImportError:
    HAS_TAXII_CLIENT = False

try:
    from stix2 import parse as stix_parse
    HAS_STIX2 = True
except ImportError:
    HAS_STIX2 = False

try:
    import requests
except ImportError:
    requests = None


def connect_taxii_server(url, username=None, password=None, version="2.1"):
    """Connect to a TAXII server and return Server object."""
    if not HAS_TAXII_CLIENT:
        print("[!] taxii2-client required: pip install taxii2-client", file=sys.stderr)
        sys.exit(1)

    kwargs = {}
    if username and password:
        kwargs["user"] = username
        kwargs["password"] = password

    print(f"[*] Connecting to TAXII {version} server: {url}")
    if version == "2.0":
        server = Server20(url, **kwargs)
    else:
        server = Server21(url, **kwargs)

    print(f"[+] Connected: {server.title or 'Untitled'}")
    return server


def discover_collections(server, version="2.1"):
    """Discover available collections on the TAXII server."""
    collections = []
    for api_root in server.api_roots:
        print(f"[*] API Root: {api_root.title or api_root.url}")
        for col in api_root.collections:
            col_info = {
                "id": col.id,
                "title": col.title or "Untitled",
                "description": getattr(col, "description", ""),
                "can_read": getattr(col, "can_read", True),
                "can_write": getattr(col, "can_write", False),
                "media_types": getattr(col, "media_types", []),
                "api_root": api_root.title or str(api_root.url),
            }
            collections.append(col_info)
            print(f"    Collection: {col_info['title']} ({col.id})")
    return collections


def fetch_collection_objects(collection, added_after=None, limit=100, obj_type=None):
    """Fetch STIX objects from a TAXII collection."""
    kwargs = {}
    if added_after:
        kwargs["added_after"] = added_after

    print(f"[*] Fetching objects from collection: {collection.title or collection.id}")
    try:
        envelope = collection.get_objects(**kwargs)
    except Exception as e:
        print(f"[!] Error fetching objects: {e}", file=sys.stderr)
        return []

    if isinstance(envelope, dict):
        objects = envelope.get("objects", [])
    elif hasattr(envelope, "objects"):
        objects = envelope.objects or []
    else:
        objects = []

    if obj_type:
        objects = [o for o in objects if o.get("type") == obj_type]

    if limit and len(objects) > limit:
        objects = objects[:limit]

    print(f"[+] Retrieved {len(objects)} object(s)")
    return objects


def extract_indicators(stix_objects):
    """Extract indicators of compromise from STIX objects."""
    indicators = []
    for obj in stix_objects:
        obj_type = obj.get("type", "")
        if obj_type == "indicator":
            indicators.append({
                "type": "indicator",
                "id": obj.get("id", ""),
                "name": obj.get("name", ""),
                "description": obj.get("description", "")[:200],
                "pattern": obj.get("pattern", ""),
                "pattern_type": obj.get("pattern_type", "stix"),
                "valid_from": obj.get("valid_from", ""),
                "valid_until": obj.get("valid_until", ""),
                "labels": obj.get("labels", []),
                "confidence": obj.get("confidence", 0),
                "created": obj.get("created", ""),
            })
    return indicators


def extract_threat_actors(stix_objects):
    """Extract threat actor information from STIX objects."""
    actors = []
    for obj in stix_objects:
        if obj.get("type") == "threat-actor":
            actors.append({
                "type": "threat-actor",
                "id": obj.get("id", ""),
                "name": obj.get("name", ""),
                "description": obj.get("description", "")[:200],
                "aliases": obj.get("aliases", []),
                "roles": obj.get("roles", []),
                "goals": obj.get("goals", []),
                "sophistication": obj.get("sophistication", ""),
                "resource_level": obj.get("resource_level", ""),
                "primary_motivation": obj.get("primary_motivation", ""),
            })
    return actors


def extract_malware(stix_objects):
    """Extract malware family info from STIX objects."""
    malware = []
    for obj in stix_objects:
        if obj.get("type") == "malware":
            malware.append({
                "type": "malware",
                "id": obj.get("id", ""),
                "name": obj.get("name", ""),
                "description": obj.get("description", "")[:200],
                "malware_types": obj.get("malware_types", []),
                "is_family": obj.get("is_family", False),
                "aliases": obj.get("aliases", []),
                "capabilities": obj.get("capabilities", []),
            })
    return malware


def extract_attack_patterns(stix_objects):
    """Extract MITRE ATT&CK patterns from STIX objects."""
    patterns = []
    for obj in stix_objects:
        if obj.get("type") == "attack-pattern":
            external_refs = obj.get("external_references", [])
            mitre_id = ""
            for ref in external_refs:
                if ref.get("source_name") in ("mitre-attack", "mitre-mobile-attack"):
                    mitre_id = ref.get("external_id", "")
                    break
            patterns.append({
                "type": "attack-pattern",
                "id": obj.get("id", ""),
                "name": obj.get("name", ""),
                "mitre_id": mitre_id,
                "description": obj.get("description", "")[:200],
                "kill_chain_phases": obj.get("kill_chain_phases", []),
            })
    return patterns


def summarize_stix_objects(stix_objects):
    """Group and count STIX objects by type."""
    type_counts = {}
    for obj in stix_objects:
        t = obj.get("type", "unknown")
        type_counts[t] = type_counts.get(t, 0) + 1
    return type_counts


def format_summary(type_counts, indicators, actors, malware, attack_patterns):
    """Print human-readable summary."""
    print(f"\n{'='*60}")
    print(f"  STIX/TAXII Feed Intelligence Report")
    print(f"{'='*60}")

    print(f"\n  Object Type Distribution:")
    for obj_type, count in sorted(type_counts.items(), key=lambda x: -x[1]):
        print(f"    {obj_type:30s}: {count}")

    if indicators:
        print(f"\n  Indicators ({len(indicators)}):")
        for ind in indicators[:10]:
            print(f"    {ind['name'][:40]:40s} | {ind['pattern_type']:8s} | "
                  f"{ind['pattern'][:50]}")

    if actors:
        print(f"\n  Threat Actors ({len(actors)}):")
        for actor in actors[:10]:
            aliases = ", ".join(actor["aliases"][:3]) if actor["aliases"] else ""
            print(f"    {actor['name']:30s} | {actor['sophistication']:12s} | {aliases}")

    if malware:
        print(f"\n  Malware Families ({len(malware)}):")
        for m in malware[:10]:
            types = ", ".join(m["malware_types"][:3]) if m["malware_types"] else ""
            print(f"    {m['name']:30s} | {types}")

    if attack_patterns:
        print(f"\n  ATT&CK Patterns ({len(attack_patterns)}):")
        for ap in attack_patterns[:10]:
            print(f"    {ap.get('mitre_id', 'N/A'):12s} | {ap['name']}")


def main():
    parser = argparse.ArgumentParser(
        description="STIX/TAXII threat intelligence feed integration agent"
    )
    parser.add_argument("--server", required=True, help="TAXII server discovery URL")
    parser.add_argument("--username", help="TAXII authentication username")
    parser.add_argument("--password", help="TAXII authentication password")
    parser.add_argument("--version", choices=["2.0", "2.1"], default="2.1",
                        help="TAXII version (default: 2.1)")
    parser.add_argument("--collection-id", help="Specific collection ID to fetch")
    parser.add_argument("--added-after", help="Only fetch objects added after date (ISO format)")
    parser.add_argument("--type", dest="obj_type",
                        help="Filter by STIX object type (e.g., indicator, malware)")
    parser.add_argument("--limit", type=int, default=500,
                        help="Max objects to retrieve (default: 500)")
    parser.add_argument("--discover-only", action="store_true",
                        help="Only discover collections, don't fetch objects")
    parser.add_argument("--output", "-o", help="Output JSON report path")
    parser.add_argument("--verbose", "-v", action="store_true")
    args = parser.parse_args()

    server = connect_taxii_server(args.server, args.username, args.password, args.version)
    collections = discover_collections(server, args.version)

    if args.discover_only:
        report = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "tool": "STIX/TAXII Client",
            "server": args.server,
            "collections": collections,
        }
        if args.output:
            with open(args.output, "w") as f:
                json.dump(report, f, indent=2)
            print(f"\n[+] Report saved to {args.output}")
        else:
            print(json.dumps(report, indent=2))
        return

    # Fetch objects from specified or all readable collections
    all_objects = []
    for col_info in collections:
        if args.collection_id and col_info["id"] != args.collection_id:
            continue
        if not col_info.get("can_read", True):
            continue
        try:
            if args.version == "2.0":
                col = Collection20(
                    f"{args.server.rstrip('/')}/collections/{col_info['id']}/",
                    user=args.username, password=args.password
                )
            else:
                col = Collection21(
                    f"{args.server.rstrip('/')}/collections/{col_info['id']}/",
                    user=args.username, password=args.password
                )
            objects = fetch_collection_objects(col, args.added_after, args.limit, args.obj_type)
            all_objects.extend(objects)
        except Exception as e:
            print(f"[!] Error fetching {col_info['title']}: {e}")

    type_counts = summarize_stix_objects(all_objects)
    indicators = extract_indicators(all_objects)
    actors = extract_threat_actors(all_objects)
    malware_items = extract_malware(all_objects)
    attack_patterns = extract_attack_patterns(all_objects)

    format_summary(type_counts, indicators, actors, malware_items, attack_patterns)

    report = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "tool": "STIX/TAXII Client",
        "server": args.server,
        "taxii_version": args.version,
        "collections_discovered": len(collections),
        "total_objects": len(all_objects),
        "type_distribution": type_counts,
        "indicators": indicators,
        "threat_actors": actors,
        "malware": malware_items,
        "attack_patterns": attack_patterns,
    }

    if args.output:
        with open(args.output, "w") as f:
            json.dump(report, f, indent=2)
        print(f"\n[+] Report saved to {args.output}")
    elif args.verbose:
        print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
STIX/TAXII 2.1 Feed Integration Script

Implements a TAXII 2.1 feed consumer that:
- Discovers TAXII server API roots and collections
- Polls collections for new STIX objects
- Parses and categorizes STIX 2.1 objects
- Extracts actionable IOCs from indicators
- Exports to multiple formats (JSON, CSV, STIX bundle)

Requirements:
    pip install taxii2-client stix2 requests

Usage:
    python process.py --server https://cti-taxii.mitre.org/taxii2/ --discover
    python process.py --server https://cti-taxii.mitre.org/taxii2/ --collection COLLECTION_ID --poll
    python process.py --server https://cti-taxii.mitre.org/taxii2/ --collection COLLECTION_ID --extract-iocs
"""

import argparse
import json
import csv
import sys
import re
from datetime import datetime, timedelta
from collections import defaultdict
from typing import Optional

try:
    from taxii2client.v21 import Server, Collection, as_pages
except ImportError:
    print("ERROR: taxii2-client not installed. Run: pip install taxii2-client")
    sys.exit(1)

try:
    from stix2 import parse, MemoryStore, Filter
except ImportError:
    print("ERROR: stix2 not installed. Run: pip install stix2")
    sys.exit(1)


class STIXTAXIIIntegrator:
    """STIX/TAXII 2.1 feed consumer and IOC extractor."""

    def __init__(self, server_url: str, user: str = "", password: str = ""):
        self.server_url = server_url
        self.user = user
        self.password = password
        self.server = None
        self.stats = defaultdict(int)

    def discover(self) -> dict:
        """Discover TAXII server API roots and collections."""
        self.server = Server(self.server_url, user=self.user, password=self.password)

        discovery = {
            "title": self.server.title,
            "description": getattr(self.server, "description", ""),
            "api_roots": [],
        }

        for api_root in self.server.api_roots:
            root_info = {
                "title": api_root.title,
                "url": api_root.url,
                "collections": [],
            }

            for collection in api_root.collections:
                root_info["collections"].append({
                    "id": collection.id,
                    "title": collection.title,
                    "description": getattr(collection, "description", ""),
                    "can_read": collection.can_read,
                    "can_write": collection.can_write,
                })
                self.stats["collections"] += 1

            discovery["api_roots"].append(root_info)

        self.stats["api_roots"] = len(discovery["api_roots"])
        return discovery

    def poll_collection(self, collection_url: str,
                        added_after: Optional[str] = None,
                        max_objects: int = 1000) -> list:
        """Poll a TAXII collection for STIX objects."""
        collection = Collection(
            collection_url, user=self.user, password=self.password
        )

        all_objects = []
        kwargs = {}
        if added_after:
            kwargs["added_after"] = added_after

        try:
            for envelope in as_pages(
                collection.get_objects, per_request=100, **kwargs
            ):
                objects = envelope.get("objects", [])
                all_objects.extend(objects)
                self.stats["objects_fetched"] += len(objects)

                if len(all_objects) >= max_objects:
                    all_objects = all_objects[:max_objects]
                    break

        except Exception as e:
            print(f"[-] Polling error: {e}")

        # Categorize by type
        for obj in all_objects:
            self.stats[f"type_{obj.get('type', 'unknown')}"] += 1

        return all_objects

    def extract_indicators(self, objects: list) -> list:
        """Extract indicator patterns from STIX objects."""
        indicators = []

        for obj in objects:
            if obj.get("type") != "indicator":
                continue

            pattern = obj.get("pattern", "")
            indicator = {
                "id": obj.get("id", ""),
                "name": obj.get("name", ""),
                "pattern": pattern,
                "pattern_type": obj.get("pattern_type", ""),
                "valid_from": obj.get("valid_from", ""),
                "valid_until": obj.get("valid_until", ""),
                "confidence": obj.get("confidence", 0),
                "indicator_types": obj.get("indicator_types", []),
                "labels": obj.get("labels", []),
                "created": obj.get("created", ""),
                "modified": obj.get("modified", ""),
            }

            # Parse pattern to extract observable values
            parsed = self._parse_stix_pattern(pattern)
            indicator["parsed_observables"] = parsed

            indicators.append(indicator)
            self.stats["indicators_extracted"] += 1

        return indicators

    def _parse_stix_pattern(self, pattern: str) -> list:
        """Parse STIX pattern to extract observable values."""
        observables = []

        # Match common STIX patterns
        ip_pattern = re.compile(
            r"ipv[46]-addr:value\s*=\s*'([^']+)'"
        )
        domain_pattern = re.compile(
            r"domain-name:value\s*=\s*'([^']+)'"
        )
        url_pattern = re.compile(
            r"url:value\s*=\s*'([^']+)'"
        )
        hash_pattern = re.compile(
            r"file:hashes\.'([^']+)'\s*=\s*'([^']+)'"
        )
        email_pattern = re.compile(
            r"email-addr:value\s*=\s*'([^']+)'"
        )

        for match in ip_pattern.finditer(pattern):
            observables.append({"type": "ip", "value": match.group(1)})

        for match in domain_pattern.finditer(pattern):
            observables.append({"type": "domain", "value": match.group(1)})

        for match in url_pattern.finditer(pattern):
            observables.append({"type": "url", "value": match.group(1)})

        for match in hash_pattern.finditer(pattern):
            observables.append({
                "type": f"hash-{match.group(1).lower()}",
                "value": match.group(2),
            })

        for match in email_pattern.finditer(pattern):
            observables.append({"type": "email", "value": match.group(1)})

        return observables

    def build_relationship_graph(self, objects: list) -> dict:
        """Build a relationship graph from STIX objects."""
        store = MemoryStore(stix_data=objects)
        relationships = store.query([Filter("type", "=", "relationship")])

        graph = {"nodes": {}, "edges": []}

        for obj in objects:
            if obj.get("type") not in ("relationship", "marking-definition"):
                graph["nodes"][obj.get("id")] = {
                    "type": obj.get("type"),
                    "name": obj.get("name", obj.get("value", obj.get("id"))),
                }

        for rel in relationships:
            graph["edges"].append({
                "source": rel.get("source_ref"),
                "target": rel.get("target_ref"),
                "type": rel.get("relationship_type"),
            })
            self.stats["relationships"] += 1

        return graph

    def export_iocs_csv(self, indicators: list, output_path: str):
        """Export parsed IOCs to CSV."""
        rows = []
        for ind in indicators:
            for obs in ind.get("parsed_observables", []):
                rows.append({
                    "indicator_id": ind["id"],
                    "indicator_name": ind["name"],
                    "observable_type": obs["type"],
                    "observable_value": obs["value"],
                    "confidence": ind["confidence"],
                    "valid_from": ind["valid_from"],
                    "valid_until": ind["valid_until"],
                })

        if rows:
            with open(output_path, "w", newline="", encoding="utf-8") as f:
                writer = csv.DictWriter(f, fieldnames=rows[0].keys())
                writer.writeheader()
                writer.writerows(rows)
            print(f"[+] Exported {len(rows)} IOCs to {output_path}")
        else:
            print("[-] No IOCs to export")

    def print_stats(self):
        """Print integration statistics."""
        print("\n=== STIX/TAXII Integration Statistics ===")
        for key, value in sorted(self.stats.items()):
            print(f"  {key.replace('_', ' ').title()}: {value}")
        print("=========================================\n")


def main():
    parser = argparse.ArgumentParser(
        description="STIX/TAXII 2.1 Feed Integration Tool"
    )
    parser.add_argument("--server", required=True, help="TAXII server URL")
    parser.add_argument("--user", default="", help="TAXII username")
    parser.add_argument("--password", default="", help="TAXII password")
    parser.add_argument("--discover", action="store_true", help="Discover server")
    parser.add_argument("--collection", help="Collection ID or URL to poll")
    parser.add_argument("--poll", action="store_true", help="Poll for objects")
    parser.add_argument("--extract-iocs", action="store_true", help="Extract IOCs")
    parser.add_argument("--added-after", help="Filter: added after timestamp")
    parser.add_argument("--max-objects", type=int, default=1000, help="Max objects")
    parser.add_argument("--output", default="stix_output.json", help="Output file")
    parser.add_argument("--csv-output", help="CSV output for IOCs")

    args = parser.parse_args()
    integrator = STIXTAXIIIntegrator(args.server, args.user, args.password)

    if args.discover:
        discovery = integrator.discover()
        print(json.dumps(discovery, indent=2))
        with open(args.output, "w") as f:
            json.dump(discovery, f, indent=2)

    elif args.collection and (args.poll or args.extract_iocs):
        # Build collection URL
        if args.collection.startswith("http"):
            collection_url = args.collection
        else:
            collection_url = (
                f"{args.server.rstrip('/')}/stix/collections/{args.collection}/"
            )

        objects = integrator.poll_collection(
            collection_url,
            added_after=args.added_after,
            max_objects=args.max_objects,
        )
        print(f"[+] Fetched {len(objects)} objects")

        if args.extract_iocs:
            indicators = integrator.extract_indicators(objects)
            print(f"[+] Extracted {len(indicators)} indicators")

            if args.csv_output:
                integrator.export_iocs_csv(indicators, args.csv_output)

            with open(args.output, "w") as f:
                json.dump(indicators, f, indent=2, default=str)
        else:
            with open(args.output, "w") as f:
                json.dump(objects, f, indent=2, default=str)

    integrator.print_stats()


if __name__ == "__main__":
    main()