Create app.py

app.py

@@ -0,0 +1,159 @@
+from crewai import Agent, Task, Crew
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+import requests
+
+# Initialize FastAPI app
+app = FastAPI()
+
+# Define agents using CrewAI
+class ThreatIntelligenceCrew:
+    def __init__(self, api_key):
+        self.api_key = api_key
+
+        # Define agents
+        self.data_collector = Agent(
+            role="Data Collector",
+            goal="Fetch threat data from external sources",
+            backstory="Specializes in collecting IOCs from threat intelligence feeds.",
+            tools=[self.fetch_iocs]  # Custom tool for fetching IOCs
+        )
+
+        self.analyst = Agent(
+            role="Threat Analyst",
+            goal="Analyze collected data for suspicious activity",
+            backstory="Expert in identifying patterns and anomalies in threat data.",
+            tools=[self.analyze_iocs]  # Custom tool for analyzing IOCs
+        )
+
+        self.correlator = Agent(
+            role="Threat Correlator",
+            goal="Correlate IOCs with known threat actors",
+            backstory="Specializes in linking IOCs to advanced threat groups.",
+            tools=[self.correlate_threats]  # Custom tool for correlation
+        )
+
+        self.reporter = Agent(
+            role="Threat Reporter",
+            goal="Generate actionable threat intelligence reports",
+            backstory="Expert in creating clear and concise reports for security teams.",
+            tools=[self.generate_report]  # Custom tool for reporting
+        )
+
+        self.responder = Agent(
+            role="Response Advisor",
+            goal="Recommend mitigation actions based on threats",
+            backstory="Specializes in providing actionable recommendations to mitigate risks.",
+            tools=[self.recommend_actions]  # Custom tool for response
+        )
+
+    # Custom tool: Fetch IOCs from AlienVault OTX
+    def fetch_iocs(self, indicator_type="ipv4"):
+        url = f"https://otx.alienvault.com/api/v1/indicators/{indicator_type}/recent"
+        headers = {"X-OTX-API-KEY": self.api_key}
+        response = requests.get(url, headers=headers)
+        if response.status_code == 200:
+            return response.json()
+        else:
+            return {"error": "Failed to fetch IOCs"}
+
+    # Custom tool: Analyze IOCs
+    def analyze_iocs(self, iocs):
+        suspicious_iocs = []
+        for ioc in iocs.get("results", []):
+            if ioc.get("pulse_info", {}).get("count", 0) > 5:  # Example threshold
+                suspicious_iocs.append(ioc)
+        return suspicious_iocs
+
+    # Custom tool: Correlate threats
+    def correlate_threats(self, iocs):
+        threat_actors = {
+            "APT28": ["1.2.3.4", "5.6.7.8"],
+            "Lazarus Group": ["9.10.11.12"]
+        }
+        correlated_threats = {}
+        for ioc in iocs:
+            ip = ioc.get("indicator")
+            for actor, ips in threat_actors.items():
+                if ip in ips:
+                    correlated_threats[ip] = actor
+        return correlated_threats
+
+    # Custom tool: Generate report
+    def generate_report(self, suspicious_iocs, correlated_threats):
+        report = {
+            "suspicious_iocs": suspicious_iocs,
+            "correlated_threats": correlated_threats,
+            "summary": f"Found {len(suspicious_iocs)} suspicious IOCs, with {len(correlated_threats)} linked to known threat actors."
+        }
+        return report
+
+    # Custom tool: Recommend actions
+    def recommend_actions(self, correlated_threats):
+        actions = []
+        for ip, actor in correlated_threats.items():
+            actions.append(f"Block IP {ip} (linked to {actor})")
+        return actions
+
+    # Define tasks for the crew
+    def create_tasks(self, indicator_type):
+        fetch_task = Task(
+            description=f"Fetch IOCs of type {indicator_type} from AlienVault OTX",
+            agent=self.data_collector,
+            expected_output="A list of IOCs in JSON format."
+        )
+
+        analyze_task = Task(
+            description="Analyze the fetched IOCs for suspicious activity",
+            agent=self.analyst,
+            expected_output="A list of suspicious IOCs."
+        )
+
+        correlate_task = Task(
+            description="Correlate suspicious IOCs with known threat actors",
+            agent=self.correlator,
+            expected_output="A dictionary mapping IOCs to threat actors."
+        )
+
+        report_task = Task(
+            description="Generate a threat intelligence report",
+            agent=self.reporter,
+            expected_output="A JSON report with suspicious IOCs, correlated threats, and a summary."
+        )
+
+        respond_task = Task(
+            description="Recommend mitigation actions based on the report",
+            agent=self.responder,
+            expected_output="A list of recommended actions."
+        )
+
+        return [fetch_task, analyze_task, correlate_task, report_task, respond_task]
+
+    # Execute the crew
+    def run_crew(self, indicator_type):
+        tasks = self.create_tasks(indicator_type)
+        crew = Crew(
+            agents=[self.data_collector, self.analyst, self.correlator, self.reporter, self.responder],
+            tasks=tasks
+        )
+        return crew.kickoff()
+
+# FastAPI endpoint
+class ThreatIntelRequest(BaseModel):
+    indicator_type: str = "ipv4"
+
+@app.post("/threat-intel")
+def threat_intel(request: ThreatIntelRequest):
+    try:
+        # Initialize the crew
+        crew = ThreatIntelligenceCrew(api_key="your_alienvault_api_key")
+        # Run the crew and get results
+        result = crew.run_crew(request.indicator_type)
+        return result
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+# Run the FastAPI app
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)