rahimizadeh commited on
Commit
d53cdbf
·
verified ·
1 Parent(s): 5da461e

Update modules/analysis.py

Browse files
Files changed (1) hide show
  1. modules/analysis.py +60 -3
modules/analysis.py CHANGED
@@ -1,8 +1,65 @@
1
  from langchain.chains import RetrievalQA
2
- from langchain_community.llms import HuggingFacePipeline # ✅ Updated import
3
  from transformers import pipeline
4
  from modules import parser, vectorizer
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5
 
6
  def run_analysis(uploaded_files, text_input, query, quick_action, temperature, start_time, end_time):
7
- print("🔥 run_analysis called") # Will appear in Application Logs
8
- return "Function was triggered!", None, None, [("INFO", "Test log processed")]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  from langchain.chains import RetrievalQA
2
+ from langchain_community.llms import HuggingFacePipeline
3
  from transformers import pipeline
4
  from modules import parser, vectorizer
5
+ from datetime import datetime
6
+ import re
7
+
8
+ def filter_logs_by_time(logs_text, start_time, end_time):
9
+ """
10
+ Filters log lines based on timestamp range.
11
+ """
12
+ if not start_time or not end_time:
13
+ return logs_text # Skip filtering if not both are set
14
+
15
+ start = datetime.fromisoformat(str(start_time))
16
+ end = datetime.fromisoformat(str(end_time))
17
+
18
+ filtered_lines = []
19
+ for line in logs_text.splitlines():
20
+ match = re.match(r"(\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2})", line)
21
+ if match:
22
+ timestamp = datetime.strptime(match.group(1), "%Y-%m-%d %H:%M:%S")
23
+ if start <= timestamp <= end:
24
+ filtered_lines.append(line)
25
+ return "\n".join(filtered_lines)
26
 
27
  def run_analysis(uploaded_files, text_input, query, quick_action, temperature, start_time, end_time):
28
+ logs_text = ""
29
+
30
+ # Combine uploaded + pasted logs
31
+ if uploaded_files:
32
+ logs_text += parser.parse_uploaded_files(uploaded_files)
33
+ if text_input:
34
+ logs_text += "\n" + text_input
35
+
36
+ if not logs_text.strip():
37
+ return "❌ No logs provided.", None, None, None
38
+
39
+ # Filter logs based on time range (if provided)
40
+ logs_text = filter_logs_by_time(logs_text, start_time, end_time)
41
+
42
+ # Use either typed query or quick action
43
+ query_text = query.strip() if query else ""
44
+ if not query_text and quick_action:
45
+ query_text = quick_action
46
+ if not query_text:
47
+ return "❌ No query or quick action selected.", None, None, None
48
+
49
+ # Process logs
50
+ docs = vectorizer.prepare_documents(logs_text)
51
+ vectordb = vectorizer.create_vectorstore(docs)
52
+
53
+ pipe = pipeline("text-generation", model="gpt2", max_length=512, temperature=temperature)
54
+ llm = HuggingFacePipeline(pipeline=pipe)
55
+
56
+ qa = RetrievalQA.from_chain_type(llm=llm, retriever=vectordb.as_retriever())
57
+ result = qa.run(query_text)
58
+
59
+ # Dummy charts and alerts for testing
60
+ bar_data = {"Hour": ["14:00", "15:00"], "Count": [8, 4]}
61
+ pie_data = {"Event Type": ["Blocked", "Scan"], "Count": [8, 4]}
62
+ alerts = [("CRITICAL", "8 blocked SSH attempts from 192.168.1.5"),
63
+ ("WARNING", "4 port scanning alerts from 10.0.0.8")]
64
+
65
+ return result, bar_data, pie_data, alerts