app.py

from ssa_code_verification import generate_response
import gradio as gr

# Define a simple function that takes an input and returns it
def identify_n_fix_vulnerability(model, user_snippet_code):
    user_message = f"""
You are expert in indentification of any code vulnerabilities under CWE and CVE standard for C, C++ programming language.
Your task is to identify any vulnerability in the snippet code that cyber attacker can manipulate the code and cause harmful to the system.
In addition, fix the code with any vulnerables found. 
The snippet code is delimited in the triple angle brakets <<<{user_snippet_code}>>>
"""
    response = ""
    #TO-DO: Check if the snippet code is actual programming language
    if len(user_snippet_code) != 0:
        response = generate_response(model, user_message)
    return response

# Define the options for the dropdown
model_options = ["llama2-7b", "llama2-13b", "llama2-70b"]

# # CCS setup for front end
# custom_css = """
#     .gr-textbox {
#         height: 1000px; 
#     }

#     .gr-dropdown {
#         height: 50px; 
#     }

#     .gr-iframe-container .title {
#         font-size: 80px;
#         color: #3366cc; 
#         font-family: 'Arial', sans-serif; 
#         font-weight: bold;
#     }
# """

title_html = """
    <h1 style='
        font-size: 80px; 
        color: #f27d52; 
        font-family: Arial, sans-serif; 
        font-weight: bold;
        background-color: #d5e6e6;
        text-align: center;
    '>
    Cyber Reasoning System</h1>
"""
# Create a Gradio interface with a text input and a dropdown menu
iface = gr.Interface(
    fn=identify_n_fix_vulnerability, 
    inputs=[
        gr.Dropdown(model_options, value = "llama2-7b", label = 'SLM Options', info="Please, select the llama2 model options"), 
        gr.Textbox(
            placeholder="Enter Snippet Code", 
            label="Code Snippet", 
            info="Please, provide snippet code for vulnerability finding and fixing",
            lines=25
        )
    ], 
    outputs= gr.Textbox(
        placeholder= " ", 
        label= "SSA Solutions",
        info= "Vulnerabilities are identified and fixed as follows",
        lines= 31
    ),
    title=title_html
)

# Launch the interface
iface.launch()