Malware_LLM / app.py
Krish30's picture
Update app.py
e298a4a verified
raw
history blame contribute delete
No virus
3.51 kB
import streamlit as st
import google.generativeai as genai
import os
import PyPDF2 as pdf
from dotenv import load_dotenv
load_dotenv()
genai.configure(api_key=("AIzaSyDziGvuT1woHnH4_S3L_zQZV55Yj-123A8"))
#alternative key
#genai.configure(api_key=("AIzaSyAr3d_7fp0wMxuUrnf_tATknu_TRPKDdxg"))
# gemini function for general content generation
def get_gemini_response(input):
model = genai.GenerativeModel('gemini-pro')
response = model.generate_content(input)
return response
# convert pdf to text
def input_pdf_text(uploaded_file):
reader = pdf.PdfReader(uploaded_file)
text = ""
for page in range(len(reader.pages)):
page = reader.pages[page]
text += str(page.extract_text())
return text
# malware detection function
def detect_malware(input_text):
malware_prompt = f"""
### As a cybersecurity expert, your task is to analyze the following text for any indications of malware.
### Text:
{input_text}
### Analysis Output:
1. Identify any potential malware-related content.
2. Explain the reasoning behind your identification.
3. Provide recommendations for mitigating any identified risks.
"""
response = get_gemini_response(malware_prompt)
return response
# chatbot function
def chatbot_response(user_input):
chatbot_prompt = f"""
### You are an intelligent and friendly chatbot. Engage in a meaningful conversation with the user.
### User Input:
{user_input}
### Chatbot Response:
"""
response = get_gemini_response(chatbot_prompt)
return response
# Function to parse and display response content
def display_response_content(response):
st.subheader("Response Output")
if response and response.candidates:
response_content = response.candidates[0].content.parts[0].text if response.candidates[0].content.parts else ""
sections = response_content.split('###')
for section in sections:
if section.strip():
section_lines = section.split('\n')
section_title = section_lines[0].strip()
section_body = '\n'.join(line.strip() for line in section_lines[1:] if line.strip())
if section_title:
st.markdown(f"**{section_title}**")
if section_body:
st.write(section_body)
else:
st.write("No response received from the model.")
## Streamlit App
st.title("AI-Powered Security and Chatbot System")
st.text("Use the AI system for malware detection and Awaring yourself.")
# Tabs for different functionalities
tab1, tab2 = st.tabs(["Malware Detection", "Chatbot"])
with tab1:
st.header("Malware Detection")
uploaded_file = st.file_uploader("Upload a file for malware detection", type="pdf", help="Please upload a PDF file.")
submit_malware = st.button('Analyze for Malware')
if submit_malware:
if uploaded_file is not None:
text = input_pdf_text(uploaded_file)
response = detect_malware(text)
# Parse and display response in a structured way
display_response_content(response)
with tab2:
st.header("Chatbot")
user_input = st.text_input("Type your message here")
submit_chat = st.button('Send')
if submit_chat:
if user_input:
response = chatbot_response(user_input)
# Parse and display response in a structured way
display_response_content(response)