opcode3/CyberSecurity_Chatbot

1 2# Import required libraries3from transformers import GPT2LMHeadModel, GPT2Tokenizer4from sentence_transformers import SentenceTransformer, util5import torch6import streamlit as st7 8# Load GPT-2 Model and Tokenizer from Hugging Face Hub9def load_gpt2_model():10    model_name = 'gpt2'  # or 'gpt2-medium', 'gpt2-large', 'gpt2-xl'11    model = GPT2LMHeadModel.from_pretrained(model_name)12    tokenizer = GPT2Tokenizer.from_pretrained(model_name)13    tokenizer.pad_token = tokenizer.eos_token  # Set the padding token14    return model, tokenizer15 16model, tokenizer = load_gpt2_model()17 18# Load Sentence Transformer Model19embedder = SentenceTransformer('all-MiniLM-L6-v2')20 21# Expanded Knowledge Base22 23# expanded_knowledge_base = [24#     "Cybersecurity involves protecting computer systems from theft or damage to hardware, software, or data.",25#     "A security breach is an incident that results in unauthorized access to computer data, applications, networks, or devices.",26#     "Phishing is a method of trying to gather personal information using deceptive e-mails and websites.",27#     "Malware is software that is specifically designed to disrupt, damage, or gain unauthorized access to a computer system.",28#     "A firewall is a network security device that monitors and controls incoming and outgoing network traffic based on predetermined security rules.",29#     "Encryption is the process of converting information or data into a code, especially to prevent unauthorized access.",30#     "Two-factor authentication (2FA) is a security process in which the user provides two different authentication factors to verify themselves.",31#     "A Distributed Denial-of-Service (DDoS) attack is an attempt to make an online service unavailable by overwhelming it with traffic from multiple sources.",32#     "Social engineering is the use of deception to manipulate individuals into divulging confidential or personal information that may be used for fraudulent purposes.",33#     "Ransomware is a type of malware that threatens to publish the victim's personal data or perpetually block access to it unless a ransom is paid.",34#     "A VPN, or Virtual Private Network, extends a private network across a public network and enables users to send and receive data as if their devices were directly connected to the private network.",35#     "Intrusion detection systems (IDS) are devices or software applications that monitor a network for malicious activity or policy violations.",36#     "Zero-day vulnerabilities are software vulnerabilities that are unknown to the software vendor and have no patch or fix.",37#     "Penetration testing, also known as pen testing, is a simulated cyber attack against your computer system to check for exploitable vulnerabilities.",38#     "A security policy is a document that outlines how to protect an organization's physical and information technology assets.",39#     "Incident response is an organized approach to addressing and managing the aftermath of a security breach or cyberattack.",40#     "Threat intelligence is information about threats and threat actors that helps mitigate harmful events in cyberspace.",41#     "SIEM (Security Information and Event Management) solutions provide real-time analysis of security alerts generated by network hardware and applications.",42#     "Access control is a method of guaranteeing that users are who they say they are and that they have the appropriate access to company data.",43#     "An attack vector is a path or means by which a hacker can gain access to a computer or network server in order to deliver a payload or malicious outcome."44#     # Add more entries45# ]46 47 48expanded_knowledge_base = [49    "Cybersecurity involves protecting computer systems from theft or damage to hardware, software, or data.",50    "A security breach is an incident that results in unauthorized access to computer data, applications, networks, or devices.",51    "Phishing is a method of trying to gather personal information using deceptive e-mails and websites.",52    "Malware is software that is specifically designed to disrupt, damage, or gain unauthorized access to a computer system.",53    "A firewall is a network security device that monitors and controls incoming and outgoing network traffic based on predetermined security rules.",54    "Encryption is the process of converting information or data into a code, especially to prevent unauthorized access.",55    "Two-factor authentication (2FA) is a security process in which the user provides two different authentication factors to verify themselves.",56    "A Distributed Denial-of-Service (DDoS) attack is an attempt to make an online service unavailable by overwhelming it with traffic from multiple sources.",57    "Social engineering is the use of deception to manipulate individuals into divulging confidential or personal information that may be used for fraudulent purposes.",58    "Ransomware is a type of malware that threatens to publish the victim's personal data or perpetually block access to it unless a ransom is paid.",59    "A VPN, or Virtual Private Network, extends a private network across a public network and enables users to send and receive data as if their devices were directly connected to the private network.",60    "Intrusion detection systems (IDS) are devices or software applications that monitor a network for malicious activity or policy violations.",61    "Zero-day vulnerabilities are software vulnerabilities that are unknown to the software vendor and have no patch or fix.",62    "Penetration testing, also known as pen testing, is a simulated cyber attack against your computer system to check for exploitable vulnerabilities.",63    "A security policy is a document that outlines how to protect an organization's physical and information technology assets.",64    "Incident response is an organized approach to addressing and managing the aftermath of a security breach or cyberattack.",65    "Threat intelligence is information about threats and threat actors that helps mitigate harmful events in cyberspace.",66    "SIEM (Security Information and Event Management) solutions provide real-time analysis of security alerts generated by network hardware and applications.",67    "Access control is a method of guaranteeing that users are who they say they are and that they have the appropriate access to company data.",68    "An attack vector is a path or means by which a hacker can gain access to a computer or network server in order to deliver a payload or malicious outcome.",69    "Access Control List (ACL) is a list of permissions attached to an object, specifying which users or system processes can access the object and what operations they can perform.",70    "Botnet Command and Control (C2) is the infrastructure used by attackers to manage and control a network of compromised devices.",71    "Container Security refers to measures and tools used to secure containerized applications and environments from threats and vulnerabilities.",72    "Data Masking is the process of obscuring specific data within a database to protect it from unauthorized access while maintaining its usability.",73    "Denial of Service (DoS) Attack is an attack designed to shut down a machine or network, making it inaccessible to intended users by overwhelming it with traffic.",74    "Endpoint Detection and Response (EDR) provides continuous monitoring and response capabilities for endpoint devices.",75    "A Security Incident is any event that compromises the confidentiality, integrity, or availability of information or information systems.",76    "Network Access Control (NAC) technologies enforce policies for accessing a network, including authentication, authorization, and compliance checks.",77    "Decryption is the process of converting encrypted data back into its original, readable format.",78    "Network Address Translation (NAT) is a method used to modify network address information in packet headers while in transit across a traffic routing device.",79    "Threat Modeling is the process of identifying potential threats and vulnerabilities in a system and designing countermeasures to mitigate them.",80    "Business Continuity Plan (BCP) is a strategy for ensuring that critical business functions continue to operate in the event of a major disruption or disaster.",81    "Disaster Recovery Plan (DRP) is a documented process or set of procedures to recover and protect a business IT infrastructure in the event of a disaster.",82    "Penetration Testing Tools are software tools used by security professionals to simulate attacks and assess the security of a system or network (e.g., Metasploit, Burp Suite).",83    "Incident Command System (ICS) is a standardized approach to the command, control, and coordination of emergency response that integrates with incident response processes.",84    "A Security Patch is an update to software or hardware designed to fix security vulnerabilities or bugs.",85    "Security Token Service (STS) is a service that issues security tokens used to access resources or services in a secure manner.",86    "Patch Tuesday is the second Tuesday of each month when Microsoft releases security patches and updates for its products.",87    "Cybersecurity Insurance is a type of insurance coverage designed to help organizations mitigate financial losses from cyber incidents or breaches.",88    "Threat Landscape is the evolving environment of potential threats and vulnerabilities that organizations must navigate to protect their assets and information."89]90 91 92# Update Embeddings for the Expanded Knowledge Base93expanded_knowledge_embeddings = embedder.encode(expanded_knowledge_base, convert_to_tensor=True)94 95# Implement Retrieval-Augmented Generation96def retrieve_knowledge(query, knowledge_base, knowledge_embeddings, top_k=1):97    query_embedding = embedder.encode(query, convert_to_tensor=True)98    hits = util.semantic_search(query_embedding, knowledge_embeddings, top_k=top_k)99    relevant_info = knowledge_base[hits[0][0]['corpus_id']]100    return relevant_info101 102def generate_response_with_knowledge(query, model, tokenizer, knowledge_base, knowledge_embeddings):103    relevant_info = retrieve_knowledge(query, knowledge_base, knowledge_embeddings)104    combined_input = f"Relevant information: {relevant_info}\nResponse:"105    #\nUser query: {query}106    inputs = tokenizer(combined_input, return_tensors='pt')107    attention_mask = inputs.attention_mask108    outputs = model.generate(109        inputs.input_ids,110        attention_mask=attention_mask,111        max_length=100,  # Limit the max length of the generated text112        num_return_sequences=1,113        temperature=0.7,114        top_p=0.9,115        do_sample=True,  # Enable sampling116        pad_token_id=tokenizer.eos_token_id117    )118    response = tokenizer.decode(outputs[0], skip_special_tokens=True)119    return response.split("Response:")[1].strip()120 121# Create the Streamlit App122def main():123    st.title("Cybersecurity Risk Assessment Chatbot")124 125    # Initialize or get the conversation history from session state126    if 'history' not in st.session_state:127        st.session_state.history = []128 129    # Display conversation history130    for query, response in st.session_state.history:131        st.write(f"You: {query}")132        st.write(f"Chatbot: {response}")133 134    # Input for the new query135    user_input = st.text_input("You:", "")136 137    if st.button("Generate Response"):138        if user_input:139            response = generate_response_with_knowledge(user_input, model, tokenizer, expanded_knowledge_base, expanded_knowledge_embeddings)140            # Append new query and response to the history141            st.session_state.history.append((user_input, response))142            # Display the new response143            st.write(f"You: {user_input}")144            st.write(f"Chatbot: {response}")145        else:146            st.write("Please enter a query to get a response.")147 148if __name__ == "__main__":149    main()150 151