Update app.py

app.py

@@ -1,159 +1,41 @@
-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)
+import streamlit as st
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+# Load the model and tokenizer
+model_name = "rajrakeshdr/IntelliSoc"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
+# Streamlit app title
+st.title("IntelliSoc Text Generation")
+
+# Input prompt
+prompt = st.text_area("Enter your prompt:", "Once upon a time")
+
+# Slider for max length
+max_length = st.slider("Max length of generated text", 50, 200, 100)
+
+# Generate text on button click
+if st.button("Generate Text"):
+    # Tokenize input
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True)
+    
+    # Generate text
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs.input_ids,
+            max_length=max_length,
+            num_return_sequences=1,
+            no_repeat_ngram_size=2,
+            top_k=50,
+            top_p=0.95,
+            temperature=0.7
+        )
+    
+    # Decode the generated text
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Display the generated text
+    st.write("Generated Text:")
+    st.write(generated_text)