epowell101
commited on
Commit
·
ac8bc57
1
Parent(s):
9ae3d6f
Added Streamlit app code and requirements file
Browse files- HF_embed_mitre_streamlit.py +102 -0
- requirements.txt +3 -0
HF_embed_mitre_streamlit.py
ADDED
|
@@ -0,0 +1,102 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import streamlit as st
|
| 2 |
+
import requests
|
| 3 |
+
import csv
|
| 4 |
+
from io import StringIO
|
| 5 |
+
|
| 6 |
+
# Required NetFlow schema
|
| 7 |
+
required_columns = [
|
| 8 |
+
'Flow duration', 'Source port', 'Destination port',
|
| 9 |
+
'Total forward packets', 'Total backward packets',
|
| 10 |
+
'Avg forward segment size', 'Avg backward segment size'
|
| 11 |
+
]
|
| 12 |
+
|
| 13 |
+
# Streamlit UI
|
| 14 |
+
st.title("NetFlow Log Comparison Tool")
|
| 15 |
+
st.write("Compare your NetFlow logs against Sigma rules or MITRE ATT&CK patterns using RAG.")
|
| 16 |
+
|
| 17 |
+
# Instructions for data upload
|
| 18 |
+
st.markdown("""
|
| 19 |
+
**Instructions:**
|
| 20 |
+
- Upload a CSV file with your NetFlow log data.
|
| 21 |
+
- Ensure that the file contains **all the required columns** listed below.
|
| 22 |
+
- You can upload **up to 5 rows** for analysis.
|
| 23 |
+
""")
|
| 24 |
+
|
| 25 |
+
# Display required schema for users
|
| 26 |
+
st.write("### Required NetFlow Schema:")
|
| 27 |
+
st.write(", ".join(required_columns))
|
| 28 |
+
|
| 29 |
+
# Step 1: File Upload
|
| 30 |
+
uploaded_file = st.file_uploader("Upload your NetFlow log sequence CSV file", type="csv")
|
| 31 |
+
|
| 32 |
+
# Step 2: User Token Input
|
| 33 |
+
hugging_face_api_token = st.text_input("Enter your Hugging Face API Token", type="password")
|
| 34 |
+
if not hugging_face_api_token:
|
| 35 |
+
st.warning("Please provide a Hugging Face API Token to proceed.")
|
| 36 |
+
|
| 37 |
+
# Step 3: Run Comparison if File Uploaded and Token Provided
|
| 38 |
+
if uploaded_file and hugging_face_api_token:
|
| 39 |
+
# Read and display the file using CSV module
|
| 40 |
+
csv_file = StringIO(uploaded_file.getvalue().decode("utf-8"))
|
| 41 |
+
reader = csv.DictReader(csv_file)
|
| 42 |
+
csv_data = list(reader)
|
| 43 |
+
|
| 44 |
+
# Display a few rows to the user
|
| 45 |
+
st.write("Uploaded File:")
|
| 46 |
+
for i, row in enumerate(csv_data[:5]):
|
| 47 |
+
st.write(row)
|
| 48 |
+
|
| 49 |
+
# Check if the file has the required schema
|
| 50 |
+
if all(col in reader.fieldnames for col in required_columns):
|
| 51 |
+
if len(csv_data) <= 5:
|
| 52 |
+
st.success("File contains all required columns and meets the row limit of 5.")
|
| 53 |
+
|
| 54 |
+
# Prepare data for Hugging Face API call
|
| 55 |
+
input_texts = [f"{row}" for row in csv_data] # Convert each row to a string for comparison
|
| 56 |
+
|
| 57 |
+
# Step 4: Call Hugging Face API
|
| 58 |
+
HUGGING_FACE_API_URL = "https://api-inference.huggingface.co/models/sentence-transformers/all-distilroberta-v1"
|
| 59 |
+
headers = {"Authorization": f"Bearer {hugging_face_api_token}"}
|
| 60 |
+
|
| 61 |
+
try:
|
| 62 |
+
# Perform inference using Hugging Face API
|
| 63 |
+
response = requests.post(HUGGING_FACE_API_URL, headers=headers, json={"inputs": input_texts})
|
| 64 |
+
response.raise_for_status()
|
| 65 |
+
|
| 66 |
+
# Display the results
|
| 67 |
+
st.write("### Comparison Results")
|
| 68 |
+
comparison_results = response.json()
|
| 69 |
+
st.write(comparison_results)
|
| 70 |
+
|
| 71 |
+
except requests.exceptions.RequestException as e:
|
| 72 |
+
st.error(f"Error calling Hugging Face API: {str(e)}")
|
| 73 |
+
|
| 74 |
+
else:
|
| 75 |
+
st.error(f"File exceeds the row limit of 5. Your file contains {len(csv_data)} rows.")
|
| 76 |
+
else:
|
| 77 |
+
missing_columns = [col for col in required_columns if col not in reader.fieldnames]
|
| 78 |
+
st.error(f"Missing columns: {', '.join(missing_columns)}")
|
| 79 |
+
|
| 80 |
+
# Step 5: Survey Link
|
| 81 |
+
st.write("### Feedback Survey")
|
| 82 |
+
st.write("We value your feedback. [Fill out our survey](https://docs.google.com/forms/d/1-P_7Uv5OphSWhTyoPuO0jjUQnYg_Hv5oVGBkhbg-H8g/prefill)") # Replace with your survey link
|
| 83 |
+
|
| 84 |
+
# Footer
|
| 85 |
+
st.markdown("---")
|
| 86 |
+
st.write("This free site is maintained by DeepTempo.")
|
| 87 |
+
st.image(".streamlit/Final DeepTempo logo.png", width=300) # Adjust the path and width as needed
|
| 88 |
+
st.write("[Visit DeepTempo.ai](https://deeptempo.ai)")
|
| 89 |
+
st.write("[Check out the underlying code on GitHub](https://github.com/deepsecoss)")
|
| 90 |
+
|
| 91 |
+
# CSS to change link color to white
|
| 92 |
+
st.markdown(
|
| 93 |
+
"""
|
| 94 |
+
<style>
|
| 95 |
+
a {
|
| 96 |
+
color: white !important;
|
| 97 |
+
text-decoration: underline; /* Optional: to keep the link recognizable */
|
| 98 |
+
}
|
| 99 |
+
</style>
|
| 100 |
+
""",
|
| 101 |
+
unsafe_allow_html=True
|
| 102 |
+
)
|
requirements.txt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
streamlit==1.25.0
|
| 2 |
+
boto3
|
| 3 |
+
requests
|