import gradio as gr
from mcp_server import *

FANCY_OUTPUTS = [
    "🔐 Calculate Hash", "🧵 Extract Strings", "📊 Calculate Entropy", 
    "🧬 Match Yara Rules", "🔍 Run Capa Analysis", 
    "🛡️ Get Results from 3rd party antivirus", "🧪 Get Sandbox Results"
]

OUTPUTS = [
    "Calculate Hash", "Extract Strings", "Calculate Entropy", 
    "Match Yara Rules", "Run Capa Analysis", 
    "Get Results from 3rd party antivirus", "Get Sandbox Results"
]

def handle_file_upload(file, checked_features, email_address=None):
    if file is None:
        return "No file uploaded."

    res = {}
    file_hash = get_file_hash(file.name)
    for i in OUTPUTS:
        if i in checked_features:
            if i == "Calculate Hash":
                res[i] = get_file_hash(file.name)
            elif i == "Extract Strings":
                res[i] = extract_strings(file.name)
            elif i == "Calculate Entropy":
                res[i] = file_entropy(file.name)
            elif i == "Match Yara Rules":
                res[i] = run_compiled_yara(file.name)
            elif i == "Get Results from 3rd party antivirus":
                res[i] = get_antivirus_detailed_reports(file_hash)
            elif i == "Get Sandbox Results":
                res[i] = get_sandbox_detailed_reports(file_hash)
            elif i == "Run Capa Analysis":
                res[i] = capa_malware_analysis(file.name)
        else:
            res[i] = f"{i} not selected."
        
    if email_address:
        try:
            send_email(email_address, str(res), f"Malware Analysis Results for {file.name}")
            email_status = "Email sent successfully."
        except Exception as e:
            print(f"Error sending email: {str(e)}")
            email_status = f"Failed to send email: {str(e)}"
    else:
        email_status = "Email not requested."

    return [res[x] for x in OUTPUTS] + [email_status]

def create_interface():

    with gr.Blocks() as demo:
        gr.HTML("""
        <style>

        .selected input.svelte-1e02hys{
                background-color:#0e203f!important;
                color: white !important;
                fill: white !important;
                accent-color: white !important;
                border-color: white !important;
        }

        #feature_checkbox_group {
        padding: 10px;
        border-radius: 10px;
        background-color:#0e203f; 
    }
        #static_analysis_accordion {
            background-color: #0e203f !important;
            color: black;
            border-radius: 8px;
            padding: 8px;
        }
        #capa_analysis_accordion {
            background-color: #0e203f !important; 
            color: black;
            border-radius: 8px;
            padding: 8px;
        }
        #cyber_threat_intelligence_accordion {
            background-color: #0e203f !important;
            color: black;
            border-radius: 8px;
            padding: 8px;
        }   
        #email_status_box {
            background-color: #1b3d77 !important;
            color: black;
            border-radius: 8px;
            padding: 8px;
        }
                
        #submit_button {
            background-color: #0e203f !important;
            color: white;
            border-radius: 8px;
            padding: 8px;
        }
        #submit_button:hover {
            background-color: #1b3d77 !important;
            color: white;
        }
        </style>
        """)
        gr.Markdown("# Malware Analysis Toolkit")
        gr.Image("images/header.png", height=150, show_label=False, show_download_button=False, container=False, elem_id="logo")
        gr.Markdown("Analyze files using CAPA, YARA, entropy, string extraction, and VirusTotal integrations.")
        gr.Markdown("This is created in order to be used as a MCP for malware analysis. \
                    As a result this UI is not fully functional and is meant to be installed locally in addition to \
                    an LLM that supports MCP connections. In order to run the MCP server you need to install and run \
                    it locally, for instructions please refer to the README")
        with gr.Row():
            with gr.Column(elem_id="input_column"):
                input_file = gr.File(label="File to be analysed")
                feature_checklist = gr.CheckboxGroup(
                choices=OUTPUTS,
                label="Select Analysis Features",
                elem_id="feature_checkbox_group"
                )
                send_email_checkbox = gr.Checkbox(label="Send results by email?")
                email_input = gr.Textbox(label="Email Address", visible=False)
                submit_button = gr.Button("Submit", elem_id="submit_button")
            with gr.Column():
                with gr.Accordion("Static Analysis", open=True, elem_id="static_analysis_accordion"):
                    output_hash = gr.Textbox(label=FANCY_OUTPUTS[0], interactive=False)
                    output_strings = gr.Textbox(label=FANCY_OUTPUTS[1], interactive=False)
                    output_entropy = gr.Textbox(label=FANCY_OUTPUTS[2], interactive=False)
                    output_yara = gr.Textbox(label=FANCY_OUTPUTS[3], interactive=False)
                with gr.Accordion("Capa Analysis", open=False, elem_id="capa_analysis_accordion"):
                    output_capa = gr.Textbox(label=FANCY_OUTPUTS[4], interactive=False)
                with gr.Accordion("Cyber Threat Intelligence", open=False, elem_id="cyber_threat_intelligence_accordion"):
                    output_antivirus = gr.Textbox(label=FANCY_OUTPUTS[5], interactive=False)
                    output_sandbox = gr.Textbox(label=FANCY_OUTPUTS[6], interactive=False)
                output_boxes = [output_hash, output_strings, output_entropy, output_yara, output_capa, output_antivirus, output_sandbox]
                email_status = gr.Textbox(label="Email Status", interactive=False, elem_id="email_status_box")


        send_email_checkbox.change(
            lambda checked: gr.update(visible=checked, interactive=checked),
            inputs=send_email_checkbox,
            outputs=email_input
        )
        submit_button.click(handle_file_upload, 
                            inputs=[input_file, feature_checklist, email_input], 
                            outputs=output_boxes+ [email_status])

    return demo

def run_interface():
    interface = create_interface()
    interface.launch(server_port=7860)