Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,364 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+from transformers import pipeline, AutoTokenizer, AutoModelForTokenClassification
+import torch
+import pypdf
+import pandas as pd
+import torch._dynamo
+import re
+
+# --- VIZUALIZACE ---
+from streamlit_agraph import agraph, Node, Edge, Config
+
+# --- GLIREL IMPORT ---
+from glirel import GLiREL
+
+# Potlačení chyb pro Windows
+torch._dynamo.config.suppress_errors = True
+
+st.set_page_config(page_title="CTI Intelligence Suite", page_icon="🛡️", layout="wide")
+
+# ==========================================
+# 1. NAČÍTÁNÍ MODELŮ
+# ==========================================
+
+@st.cache_resource
+def load_ner_model():
+    """ Načte SecureModernBERT pro entity. """
+    device = 0 if torch.cuda.is_available() else -1
+    model_name = "attack-vector/SecureModernBERT-NER"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForTokenClassification.from_pretrained(model_name)
+    pipe = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple", device=device)
+    return pipe
+
+@st.cache_resource
+def load_glirel_model():
+    """ Načte GLiREL pro vztahy. """
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    try:
+        model = GLiREL.from_pretrained("jackboyla/glirel-large-v0")
+        model.to(device)
+        model.eval()
+        return model
+    except Exception as e:
+        st.error(f"Nepodařilo se stáhnout model GLiREL: {e}")
+        return None
+
+# ==========================================
+# 2. LOGIKA ZPRACOVÁNÍ TEXTU
+# ==========================================
+
+def extract_text_from_pdf(uploaded_file):
+    try:
+        pdf_reader = pypdf.PdfReader(uploaded_file)
+        text = ""
+        for page in pdf_reader.pages:
+            t = page.extract_text()
+            if t: text += t + "\n\n"
+        return text
+    except Exception as e:
+        st.error(f"Error reading PDF: {e}")
+        return ""
+
+def analyze_ner_batched(pipeline, text, batch_size=8):
+    """
+    Projde CELÝ text po kusech (nezkracuje ho), aby našel všechny entity.
+    """
+    chunk_size = 4000
+    results = []
+    
+    for i in range(0, len(text), chunk_size):
+        chunk = text[i : i + chunk_size]
+        if not chunk.strip(): continue
+        
+        # Analýza kusu
+        chunk_results = pipeline(chunk)
+        
+        # Posun indexů
+        for entity in chunk_results:
+            entity['start'] += i
+            entity['end'] += i
+            results.append(entity)
+            
+    return results
+
+def merge_close_entities(results, original_text, max_char_distance=2):
+    if not results: return []
+    merged = []
+    current = results[0].copy()
+    for next_entity in results[1:]:
+        gap_start = current['end']
+        gap_end = next_entity['start']
+        if gap_start > gap_end: gap_start = gap_end
+        gap_text = original_text[gap_start:gap_end]
+        
+        if (current['entity_group'] == next_entity['entity_group'] and 
+            len(gap_text) <= max_char_distance and 
+            "." not in gap_text and "," not in gap_text):
+            current['end'] = next_entity['end']
+            current['score'] = max(current['score'], next_entity['score'])
+        else:
+            merged.append(current)
+            current = next_entity.copy()
+    merged.append(current)
+    return merged
+
+# ==========================================
+# 3. FILTRACE VĚT (SMART FILTER)
+# ==========================================
+
+def filter_text_smartly(text, ner_results):
+    """
+    Vezme celý text a výsledky NERu.
+    Vrátí pouze ty věty, které obsahují alespoň jednu entitu.
+    """
+    sentence_spans = []
+    for match in re.finditer(r'[^.!?]+[.!?]', text):
+        sentence_spans.append((match.start(), match.end(), match.group()))
+    
+    if not sentence_spans and text:
+        sentence_spans.append((0, len(text), text))
+
+    relevant_sentences = []
+    
+    for s_start, s_end, s_text in sentence_spans:
+        has_entity = False
+        for ent in ner_results:
+            if ent['start'] >= s_start and ent['end'] <= s_end:
+                has_entity = True
+                break
+        
+        if has_entity:
+            relevant_sentences.append(s_text.strip())
+
+    clean_text = " ".join(relevant_sentences)
+    # Pojistka: Limit pro GLiREL
+    return clean_text[:3000] 
+
+# ==========================================
+# 4. GLIREL LOGIKA (CLEAN & ROBUST)
+# ==========================================
+
+def align_and_predict_relations(glirel_model, text, ner_results, threshold=0.4):
+    if not glirel_model: return []
+    
+    tokens = text.split()
+    token_spans = []
+    curr = 0
+    for t in tokens:
+        start = text.find(t, curr)
+        if start == -1: start = curr
+        end = start + len(t)
+        token_spans.append((start, end))
+        curr = end
+
+    glirel_ner = []
+    token_to_full_entity = {} 
+    
+    for ent in ner_results:
+        c_start, c_end = ent['start'], ent['end']
+        full_name = text[c_start:c_end].strip()
+        t_start, t_end = -1, -1
+        for i, (ts, te) in enumerate(token_spans):
+            if ts >= c_start and t_start == -1: t_start = i
+            if te <= c_end: t_end = i
+        
+        if t_start != -1 and t_end != -1:
+            glirel_ner.append([t_start, t_end, ent['entity_group'], full_name])
+            for t_idx in range(t_start, t_end + 1):
+                token_to_full_entity[t_idx] = full_name
+
+    if not glirel_ner: return []
+
+    labels = [
+        "uses", "targets", "communicates_with", "drops", 
+        "located_at", "attributed_to", "exploits", 
+        "compromises", "downloads", "resolves_to", "variant_of"
+    ]
+
+    try:
+        relations = glirel_model.predict_relations(
+            tokens, labels, threshold=threshold, ner=glirel_ner, top_k=3
+        )
+    except Exception as e:
+        st.error(f"GLiREL Error: {e}")
+        return []
+
+    # Pomocná funkce na čištění textu (odstraní [], Unknown, prázdné stringy)
+    def clean_entity_text(raw_val):
+        if raw_val is None: return None
+        if isinstance(raw_val, list):
+            if not raw_val: return None
+            raw_val = " ".join([str(x) for x in raw_val])
+        text_val = str(raw_val).strip()
+        if text_val in ["", "[]", "['']", "Unknown", "None"]: return None
+        return text_val
+
+    best_relations = {}
+    for rel in relations:
+        head_idx = rel['head_pos'][0]
+        tail_idx = rel['tail_pos'][0]
+        
+        raw_head = token_to_full_entity.get(head_idx, rel.get('head_text'))
+        raw_tail = token_to_full_entity.get(tail_idx, rel.get('tail_text'))
+
+        head = clean_entity_text(raw_head)
+        tail = clean_entity_text(raw_tail)
+
+        if head and tail and head != tail:
+            current_score = rel['score']
+            relation_label = rel['label']
+            pair = sorted([head, tail]) 
+            unique_key = (pair[0], pair[1], relation_label)
+
+            if unique_key not in best_relations:
+                best_relations[unique_key] = {
+                    "source": head, "target": tail, "relation": relation_label, "confidence": current_score
+                }
+            else:
+                if current_score > best_relations[unique_key]['confidence']:
+                    best_relations[unique_key] = {
+                        "source": head, "target": tail, "relation": relation_label, "confidence": current_score
+                    }
+            
+    return list(best_relations.values())
+
+# ==========================================
+# 5. UI APLIKACE
+# ==========================================
+
+with st.sidebar:
+    st.title("🧭 Navigation")
+    page = st.radio("Go to:", ["Analyzer", "Visualizations"])
+    st.markdown("---")
+    
+    if page == "Analyzer":
+        st.subheader("⚙️ Settings")
+        confidence_threshold = st.slider(
+            "Relation Confidence (%)", min_value=0, max_value=100, value=59, 
+            help="Zobrazí jen vztahy, kde si je model jistý na více než X %."
+        )
+
+if page == "Analyzer":
+    st.title("CTI Analyzer")
+    
+    with st.spinner("Loading models..."):
+        ner_pipe = load_ner_model()
+
+    col1, col2 = st.columns([1, 2])
+    
+    with col1:
+        st.subheader("📂 Input")
+        uploaded_file = st.file_uploader("Upload PDF Report", type=["pdf"])
+        
+    with col2:
+        # ZMĚNA ZDE: NOVÝ DEFAULT TEXT
+        default_text = r"""Lazarus Group, often linked to the North Korean government, has been observed targeting the financial sector and cryptocurrency exchanges in Japan. The threat actor uses AppleJeus malware to infiltrate networks. The malware was found located at C:\Windows\Temp\update.exe. Security researchers attributed this campaign to Hidden Cobra. In a recent incident, Lazarus Group also targeted Sony Pictures."""        
+        if uploaded_file:
+            with st.spinner("Reading PDF..."):
+                txt = extract_text_from_pdf(uploaded_file)
+                st.info(f"Loaded {len(txt)} characters from PDF.")
+                st.text_area("Preview", value=txt[:500] + "...", height=150, disabled=True)
+        else:
+            txt = st.text_area("Or Paste Text Here", value=default_text, height=150)
+
+    st.divider()
+    if st.button("Analyze", type="primary"):
+        if not txt.strip():
+            st.warning("Please enter some text or upload a PDF.")
+        else:
+            with st.status("Running analysis...") as status:
+                
+                # 1. NER - SCAN CELÉHO DOKUMENTU
+                status.write("1. Scanning full document for Entities...")
+                full_raw_ents = analyze_ner_batched(ner_pipe, txt)
+                
+                # 2. FILTRACE TEXTU
+                status.write("2. Selecting key sentences...")
+                optimized_text = filter_text_smartly(txt, full_raw_ents)
+                if not optimized_text: optimized_text = txt[:2000]
+                
+                st.info(f"Text optimized from {len(txt)} to {len(optimized_text)} chars.")
+
+                # 3. RE-ALIGNMENT
+                status.write("3. Re-aligning entities...")
+                final_raw_ents = analyze_ner_batched(ner_pipe, optimized_text)
+                final_ents = merge_close_entities(final_raw_ents, optimized_text)
+                
+                # 4. GLIREL
+                status.write("4. Extracting Relations...")
+                glirel = load_glirel_model()
+                threshold_float = confidence_threshold / 100.0
+                rels = align_and_predict_relations(glirel, optimized_text, final_ents, threshold=threshold_float)
+                
+                status.update(label="Done!", state="complete")
+
+            # Uložení výsledků
+            df_ents = pd.DataFrame([{
+                "Entity": optimized_text[e['start']:e['end']], 
+                "Type": e['entity_group'], 
+                "Confidence": e['score']
+            } for e in final_ents])
+            
+            # Konverze NER dat na JSON string pro download
+            ner_json = df_ents.to_json(orient="records", indent=4)
+
+            # Čistění duplicit ve vztazích
+            df_rels = pd.DataFrame(rels)
+            if not df_rels.empty:
+                df_rels = df_rels.drop_duplicates(subset=["source", "target", "relation"])
+
+            st.session_state['data'] = df_ents
+            st.session_state['rels'] = df_rels
+
+            c1, c2 = st.columns(2)
+            with c1:
+                st.subheader(f"Entities ({len(df_ents)})")
+                st.dataframe(df_ents, use_container_width=True)
+                
+                # --- TLAČÍTKO PRO DOWNLOAD JSON ---
+                st.download_button(
+                    label="📥 Download NER JSON",
+                    data=ner_json,
+                    file_name="ner_entities.json",
+                    mime="application/json"
+                )
+                # ----------------------------------
+                
+            with c2:
+                st.subheader(f"Relations ({len(df_rels)})")
+                st.dataframe(st.session_state['rels'], use_container_width=True)
+
+elif page == "Visualizations":
+    st.title("Knowledge Graph")
+    if 'data' in st.session_state and not st.session_state['data'].empty:
+        nodes, edges = [], []
+        added = set()
+        
+        type_colors = {"MALWARE": "#ff4b4b", "ACTOR": "#ffa421", "TOOL": "#1c83e1", "IP": "#21c354"}
+        
+        for _, row in st.session_state['data'].iterrows():
+            ent = row['Entity']
+            if ent not in added:
+                color = type_colors.get(row['Type'], "#888")
+                nodes.append(Node(id=ent, label=ent, size=20, color=color))
+                added.add(ent)
+        
+        if 'rels' in st.session_state and not st.session_state['rels'].empty:
+            for _, row in st.session_state['rels'].iterrows():
+                if row['source'] not in added:
+                    nodes.append(Node(id=row['source'], label=row['source'], size=20, color="#888"))
+                    added.add(row['source'])
+                if row['target'] not in added:
+                    nodes.append(Node(id=row['target'], label=row['target'], size=20, color="#888"))
+                    added.add(row['target'])
+                    
+                edges.append(Edge(
+                    source=row['source'], target=row['target'], label=row['relation'], 
+                    color="red", arrows="to"
+                ))
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+        config = Config(width="100%", height=600, directed=True, physics=True)
+        agraph(nodes=nodes, edges=edges, config=config)
+    else:
+        st.warning("No data found. Please run analysis first.")