Update app.py

app.py

@@ -1,330 +1,155 @@
-import os
-import json
-import hashlib
-import time
-from datetime import datetime
 import gradio as gr
-from huggingface_hub import InferenceClient
-from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
-from reportlab.lib.styles import getSampleStyleSheet
-from reportlab.lib.units import inch
-
-# ============================================================
-# CONFIGURATION
-# ============================================================
-
-OMNISCIENT_ROOT = "/opt/omniscient"
-CASES_ROOT = os.path.join(OMNISCIENT_ROOT, "cases")
-os.makedirs(CASES_ROOT, exist_ok=True)
-
-MODEL_OPTIONS = {
-    "Zephyr 7B (Stable)": "HuggingFaceH4/zephyr-7b-beta",
-    "Mixtral 8x7B": "mistralai/Mixtral-8x7B-Instruct-v0.1",
-}
-
-FALLBACK_MODEL = "HuggingFaceH4/zephyr-7b-beta"
-
-# ============================================================
-# UTILITIES
-# ============================================================
-
-def sha256_data(data):
-    return hashlib.sha256(data.encode()).hexdigest()
-
-def sha256_file(path):
-    h = hashlib.sha256()
-    with open(path, "rb") as f:
-        while chunk := f.read(8192):
-            h.update(chunk)
-    return h.hexdigest()
-
-def case_path(case_id):
-    return os.path.join(CASES_ROOT, case_id)
-
-def manifest_path(case_id):
-    return os.path.join(case_path(case_id), "manifest.json")
-
-def audit_path(case_id):
-    return os.path.join(case_path(case_id), "audit_log.json")
-
-def ai_notes_path(case_id):
-    return os.path.join(case_path(case_id), "ai_analysis.json")
-
-# ============================================================
-# CASE MANAGEMENT
-# ============================================================
-
-def create_case(case_id):
-    path = case_path(case_id)
-    os.makedirs(path, exist_ok=True)
-    os.makedirs(os.path.join(path, "evidence"), exist_ok=True)
-    return f"✅ Case created: {case_id}"
-
-def list_cases():
-    return "\n".join(os.listdir(CASES_ROOT))
-
-# ============================================================
-# HASH-CHAINED MANIFEST
-# ============================================================
-
-def register_evidence(case_id, file, operator, notes):
-    if not file:
-        return "No file uploaded."
-
-    path = case_path(case_id)
-    evidence_dir = os.path.join(path, "evidence")
-
-    if not os.path.exists(path):
-        return "Case does not exist."
-
-    dest = os.path.join(evidence_dir, os.path.basename(file.name))
-    with open(file.name, "rb") as src, open(dest, "wb") as dst:
-        dst.write(src.read())
-
-    file_hash = sha256_file(dest)
-
-    if os.path.exists(manifest_path(case_id)):
-        with open(manifest_path(case_id), "r") as f:
-            manifest = json.load(f)
-    else:
-        manifest = []
-
-    prev_hash = manifest[-1]["entry_hash"] if manifest else "GENESIS"
-
-    entry = {
-        "file": dest,
-        "sha256": file_hash,
-        "operator": operator,
-        "notes": notes,
-        "timestamp": datetime.utcnow().isoformat(),
-        "previous_hash": prev_hash
+import requests
+from bs4 import BeautifulSoup
+from duckduckgo_search import DDGS
+from sentence_transformers import SentenceTransformer
+from transformers import pipeline
+import numpy as np
+import re
+from sklearn.metrics.pairwise import cosine_similarity
+
+# -------------------------
+# Load Models (HF Safe)
+# -------------------------
+
+embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+
+llm = pipeline(
+    "text-generation",
+    model="HuggingFaceH4/zephyr-7b-beta",
+    max_new_tokens=512,
+)
+
+# -------------------------
+# Simple In-Memory Vector Store
+# -------------------------
+
+documents = []
+embeddings = []
+
+def add_to_rag(text):
+    global documents, embeddings
+    documents.append(text)
+    embeddings.append(embedder.encode(text))
+
+def retrieve_from_rag(query, top_k=3):
+    if not embeddings:
+        return []
+    query_emb = embedder.encode(query)
+    sims = cosine_similarity([query_emb], embeddings)[0]
+    top_idx = np.argsort(sims)[-top_k:][::-1]
+    return [documents[i] for i in top_idx]
+
+# -------------------------
+# URL Scraper
+# -------------------------
+
+def scrape_url(url):
+    try:
+        r = requests.get(url, timeout=10)
+        soup = BeautifulSoup(r.text, "html.parser")
+        text = soup.get_text(separator=" ")
+        add_to_rag(text)
+        return "Scraped and added to RAG."
+    except:
+        return "Scrape failed."
+
+# -------------------------
+# DuckDuckGo Search
+# -------------------------
+
+def ddg_search(query):
+    results = []
+    with DDGS() as ddgs:
+        for r in ddgs.text(query, max_results=5):
+            results.append(r["body"])
+    combined = "\n".join(results)
+    add_to_rag(combined)
+    return combined
+
+# -------------------------
+# Hybrid Entity Extraction
+# -------------------------
+
+def regex_entities(text):
+    entities = {
+        "emails": re.findall(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}", text),
+        "phones": re.findall(r"\+?\d[\d -]{8,}\d", text),
+        "urls": re.findall(r"https?://\S+", text),
+        "ips": re.findall(r"\b(?:\d{1,3}\.){3}\d{1,3}\b", text),
     }
+    return entities
 
-    entry_hash = sha256_data(json.dumps(entry))
-    entry["entry_hash"] = entry_hash
-
-    manifest.append(entry)
-
-    with open(manifest_path(case_id), "w") as f:
-        json.dump(manifest, f, indent=2)
-
-    append_audit_log(case_id, "EVIDENCE_REGISTERED", entry)
-
-    return f"Evidence registered.\nSHA256: {file_hash}"
-
-# ============================================================
-# HASH-CHAINED AUDIT LOG
-# ============================================================
+def llm_refine_entities(text):
+    prompt = f"""
+Extract structured OSINT entities from this text.
+Return JSON with:
+people, organizations, locations, vehicles, usernames.
 
-def append_audit_log(case_id, action, details):
-    if os.path.exists(audit_path(case_id)):
-        with open(audit_path(case_id), "r") as f:
-            log = json.load(f)
-    else:
-        log = []
-
-    prev_hash = log[-1]["entry_hash"] if log else "GENESIS"
-
-    entry = {
-        "timestamp": datetime.utcnow().isoformat(),
-        "action": action,
-        "details": details,
-        "previous_hash": prev_hash
-    }
-
-    entry_hash = sha256_data(json.dumps(entry))
-    entry["entry_hash"] = entry_hash
+TEXT:
+{text}
+"""
+    output = llm(prompt)[0]["generated_text"]
+    return output
 
-    log.append(entry)
+def hybrid_extract(text):
+    base = regex_entities(text)
+    refined = llm_refine_entities(text[:2000])
+    return f"Regex Extracted:\n{base}\n\nLLM Refined:\n{refined}"
 
-    with open(audit_path(case_id), "w") as f:
-        json.dump(log, f, indent=2)
+# -------------------------
+# Chat Logic
+# -------------------------
 
-# ============================================================
-# VERIFY INTEGRITY
-# ============================================================
+def chat(query, use_web, use_rag):
+    context = ""
 
-def verify_integrity(case_id):
-    # Verify Manifest Chain
-    if os.path.exists(manifest_path(case_id)):
-        with open(manifest_path(case_id), "r") as f:
-            manifest = json.load(f)
+    if use_web:
+        context += ddg_search(query)
 
-        for i in range(1, len(manifest)):
-            if manifest[i]["previous_hash"] != manifest[i-1]["entry_hash"]:
-                return "Manifest chain broken."
+    if use_rag:
+        docs = retrieve_from_rag(query)
+        context += "\n".join(docs)
 
-    # Verify Audit Chain
-    if os.path.exists(audit_path(case_id)):
-        with open(audit_path(case_id), "r") as f:
-            log = json.load(f)
+    final_prompt = f"""
+Use the following context to answer intelligently:
 
-        for i in range(1, len(log)):
-            if log[i]["previous_hash"] != log[i-1]["entry_hash"]:
-                return "Audit chain broken."
+{context}
 
-    return "Integrity verified."
+Question: {query}
+"""
 
-# ============================================================
-# AI AUGMENTATION (CASE-BOUND)
-# ============================================================
+    response = llm(final_prompt)[0]["generated_text"]
+    return response
 
-def get_client(model_id):
-    token = os.environ.get("HF_TOKEN")
-    return InferenceClient(model=model_id, token=token)
+# -------------------------
+# Gradio UI
+# -------------------------
 
-def run_ai(model_id, messages):
-    try:
-        client = get_client(model_id)
-        response = client.chat_completion(
-            messages=messages,
-            max_tokens=1000,
-            temperature=0.4,
-            top_p=0.95,
-        )
-        return response.choices[0].message.content
-    except:
-        fallback = get_client(FALLBACK_MODEL)
-        response = fallback.chat_completion(
-            messages=messages,
-            max_tokens=1000,
-            temperature=0.4,
-            top_p=0.95,
-        )
-        return "Fallback model used.\n\n" + response.choices[0].message.content
-
-def ai_analysis(case_id, input_text, model_label):
-    system_prompt = """
-You are an AI augmentation module for digital forensics.
-Structure output with clear analytical sections.
-AI does not modify evidence.
-"""
+with gr.Blocks() as demo:
+    gr.Markdown("# 🔎 Hybrid OSINT AI Assistant")
 
-    messages = [
-        {"role": "system", "content": system_prompt},
-        {"role": "user", "content": input_text}
-    ]
+    with gr.Row():
+        query = gr.Textbox(label="Ask Question")
+        use_web = gr.Checkbox(label="Use DuckDuckGo Search")
+        use_rag = gr.Checkbox(label="Use RAG")
 
-    result = run_ai(MODEL_OPTIONS[model_label], messages)
+    chat_btn = gr.Button("Run")
 
-    # Store AI notes per case
-    notes_path = ai_notes_path(case_id)
-    if os.path.exists(notes_path):
-        with open(notes_path, "r") as f:
-            notes = json.load(f)
-    else:
-        notes = []
+    output = gr.Textbox(label="Response")
 
-    entry = {
-        "timestamp": datetime.utcnow().isoformat(),
-        "analysis": result
-    }
+    chat_btn.click(chat, inputs=[query, use_web, use_rag], outputs=output)
 
-    notes.append(entry)
+    gr.Markdown("## 🌐 URL → RAG")
+    url_input = gr.Textbox(label="Enter URL")
+    scrape_btn = gr.Button("Scrape")
+    scrape_output = gr.Textbox()
+    scrape_btn.click(scrape_url, inputs=url_input, outputs=scrape_output)
 
-    with open(notes_path, "w") as f:
-        json.dump(notes, f, indent=2)
+    gr.Markdown("## 🧩 OSINT Entity Extraction")
+    extract_input = gr.Textbox(label="Paste Text")
+    extract_btn = gr.Button("Extract Entities")
+    extract_output = gr.Textbox()
+    extract_btn.click(hybrid_extract, inputs=extract_input, outputs=extract_output)
 
-    append_audit_log(case_id, "AI_ANALYSIS", {"length": len(input_text)})
-
-    return result
-
-# ============================================================
-# CASE TIMELINE
-# ============================================================
-
-def case_timeline(case_id):
-    events = []
-
-    if os.path.exists(audit_path(case_id)):
-        with open(audit_path(case_id), "r") as f:
-            log = json.load(f)
-        for entry in log:
-            events.append(f"{entry['timestamp']} - {entry['action']}")
-
-    return "\n".join(sorted(events))
-
-# ============================================================
-# PDF EXPORT
-# ============================================================
-
-def export_case_pdf(case_id):
-    path = case_path(case_id)
-    pdf_path = os.path.join(path, f"{case_id}_report.pdf")
-
-    doc = SimpleDocTemplate(pdf_path)
-    elements = []
-    styles = getSampleStyleSheet()
-
-    elements.append(Paragraph(f"Case Report: {case_id}", styles["Heading1"]))
-    elements.append(Spacer(1, 0.3 * inch))
-
-    for file_name in ["manifest.json", "audit_log.json", "ai_analysis.json"]:
-        file_path = os.path.join(path, file_name)
-        if os.path.exists(file_path):
-            elements.append(Paragraph(file_name, styles["Heading2"]))
-            with open(file_path, "r") as f:
-                data = json.load(f)
-            elements.append(Paragraph(str(data), styles["Normal"]))
-            elements.append(Spacer(1, 0.3 * inch))
-
-    doc.build(elements)
-    return pdf_path
-
-# ============================================================
-# UI
-# ============================================================
-
-with gr.Blocks(title="Omniscient Investigative Infrastructure vNext") as demo:
-
-    gr.Markdown("## Omniscient – Advanced Case-Based Investigative Platform")
-
-    case_id = gr.Textbox(label="Case ID")
-
-    with gr.Tabs():
-
-        with gr.Tab("Case Management"):
-            create_btn = gr.Button("Create Case")
-            list_btn = gr.Button("List Cases")
-            case_output = gr.Textbox(lines=6)
-
-        with gr.Tab("Evidence"):
-            operator = gr.Textbox(label="Operator Name")
-            notes = gr.Textbox(label="Evidence Notes")
-            file_upload = gr.File(label="Upload Evidence")
-            register_btn = gr.Button("Register Evidence")
-            evidence_output = gr.Textbox(lines=10)
-
-        with gr.Tab("AI Augmentation"):
-            ai_input = gr.Textbox(lines=8)
-            model_selector = gr.Dropdown(
-                choices=list(MODEL_OPTIONS.keys()),
-                value=list(MODEL_OPTIONS.keys())[0]
-            )
-            ai_btn = gr.Button("Run AI Analysis")
-            ai_output = gr.Textbox(lines=20)
-
-        with gr.Tab("Integrity"):
-            verify_btn = gr.Button("Verify Full Integrity")
-            timeline_btn = gr.Button("View Case Timeline")
-            integrity_output = gr.Textbox(lines=15)
-
-        with gr.Tab("Export"):
-            export_btn = gr.Button("Export Case PDF")
-            export_file = gr.File()
-
-    create_btn.click(create_case, inputs=case_id, outputs=case_output)
-    list_btn.click(list_cases, outputs=case_output)
-    register_btn.click(register_evidence,
-                       inputs=[case_id, file_upload, operator, notes],
-                       outputs=evidence_output)
-    ai_btn.click(ai_analysis,
-                 inputs=[case_id, ai_input, model_selector],
-                 outputs=ai_output)
-    verify_btn.click(verify_integrity, inputs=case_id, outputs=integrity_output)
-    timeline_btn.click(case_timeline, inputs=case_id, outputs=integrity_output)
-    export_btn.click(export_case_pdf, inputs=case_id, outputs=export_file)
-
-if __name__ == "__main__":
-    demo.queue()
-    demo.launch()
+demo.launch()