import streamlit as st
import os
from dotenv import load_dotenv
import requests
import json

### Load environment variables
load_dotenv()

### Page configuration
st.set_page_config(
    page_title="🐞T&S Malware Deobfuscator",
    page_icon="🔍",
    layout="wide"
)

### Title and description
st.title("🐞 AI-Powered Malware Deobfuscator")
st.markdown("""
This tool uses AI to analyze and deobfuscate potentially malicious code.
Upload obfuscated code and get an understanding of its behaviour.
""")

### Sidebar for configuration
st.sidebar.header("⚙️ Configuration")
ai_provider = st.sidebar.radio(
    "Select AI Provider:",
    ["GitHub Models", "Azure OpenAI"]
)

### Function to call GitHub Models
def analyze_with_github_models(code, task_type):
    """
    Uses GitHub Models API to analyze code
    """
    token = os.getenv("GITHUB_TOKEN")
    
    if not token:
        return "❌ Error: GitHub token not found. Please set GITHUB_TOKEN in your environment."
    
    ### Prepare the prompt based on task
    if task_type == "deobfuscate":
        prompt = f"""You are a malware analyst. Analyze this obfuscated code and provide:
1. A deobfuscated (cleaned up, readable) version
2. Explanation of what it does
3. Potential security risks
4. If the code is too long to process, please do the following: 
 - review the code in chunks that can be processed
 - analyse strings in the top part of the code when variables are defined 
 - if this all fails, summarize the code's behaviour instead 

Do not produce an error if the obfuscated code is too long to process. Instead, follow the instructions above.

Provide clear, structured output and make comparisons with known malware patterns where applicable.

Code:
{code}"""
    elif task_type == "explain":
        prompt = f"""Explain what this code does in simple terms. Identify any malicious behavior:

Code:
{code}"""
    else:  ### yara
        prompt = f"""Generate a YARA rule to detect code similar to this:

Code:
{code}"""
    
    ### API endpoint for GitHub Models (using GPT-4o)
    url = "https://models.inference.ai.azure.com/chat/completions"
    
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {token}"
    }
    
    data = {
        "model": "gpt-4o",
        "messages": [
            {"role": "system", "content": "You are an expert malware analyst and security researcher."},
            {"role": "user", "content": prompt}
        ],
        "temperature": 0.3,
        "max_tokens": 2000
    }
    
    try:
        response = requests.post(url, headers=headers, json=data, timeout=30)
        response.raise_for_status()
        result = response.json()
        return result['choices'][0]['message']['content']
    except Exception as e:
        return f"❌ Error: {str(e)}\n\nResponse: {response.text if 'response' in locals() else 'No response'}"

### Function to call Azure OpenAI
def analyze_with_azure(code, task_type):
    """
    Uses Azure OpenAI to analyze code
    """
    endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
    api_key = os.getenv("AZURE_OPENAI_KEY")
    deployment = os.getenv("AZURE_OPENAI_DEPLOYMENT")
    
    if not all([endpoint, api_key, deployment]):
        return "❌ Error: Azure OpenAI credentials not configured."
    
    ### Prepare the prompt
    if task_type == "deobfuscate":
        prompt = f"""Analyze this obfuscated malicious code and provide:
1. Deobfuscated version
2. Explanation of functionality
3. Security threats

Code:
{code}"""
    elif task_type == "explain":
        prompt = f"""Explain this code's behavior and identify threats:\n\n{code}"""
    else:
        prompt = f"""Generate a YARA rule for this code:\n\n{code}"""
    
    url = f"{endpoint}/openai/deployments/{deployment}/chat/completions?api-version=2024-02-15-preview"
    
    headers = {
        "Content-Type": "application/json",
        "api-key": api_key
    }
    
    data = {
        "messages": [
            {"role": "system", "content": "You are an expert malware analyst."},
            {"role": "user", "content": prompt}
        ],
        "temperature": 0.3,
        "max_tokens": 2000
    }
    
    try:
        response = requests.post(url, headers=headers, json=data, timeout=30)
        response.raise_for_status()
        result = response.json()
        return result['choices'][0]['message']['content']
    except Exception as e:
        return f"❌ Error: {str(e)}"

### Main interface
col1, col2 = st.columns(2)

with col1:
    st.header("📥 Input")
    
    ### Input method selection
    input_method = st.radio("Choose input method:", ["Paste Code", "Upload File"])
    
    if input_method == "Paste Code":
        code_input = st.text_area(
            "Paste obfuscated code here:",
            height=300,
            placeholder="eval(base64_decode('...'))"
        )
    else:
        uploaded_file = st.file_uploader("Upload a file", type=['txt', 'js', 'py', 'ps1'])
        if uploaded_file:
            code_input = uploaded_file.read().decode('utf-8')
            st.text_area("File contents:", code_input, height=300)
        else:
            code_input = ""
    
    ### Analysis type
    analysis_type = st.selectbox(
        "Select analysis type:",
        ["Deobfuscate & Explain", "Quick Explanation", "Generate YARA Rule"]
    )
    
    ### Analyze button
    analyze_button = st.button("🔍 Analyze Code", type="primary")

with col2:
    st.header("📤 Results")
    
    if analyze_button:
        if not code_input:
            st.warning("⚠️ Please provide some code to analyze.")
        else:
            with st.spinner("🤖 Analyzing code..."):
                ### Map analysis type to task type
                task_map = {
                    "Deobfuscate & Explain": "deobfuscate",
                    "Quick Explanation": "explain",
                    "Generate YARA Rule": "yara"
                }
                task_type = task_map[analysis_type]
                
                ### Call appropriate API
                if ai_provider == "GitHub Models":
                    result = analyze_with_github_models(code_input, task_type)
                else:
                    result = analyze_with_azure(code_input, task_type)
                
                ### Display results
                st.markdown("#### Analysis Results")
                st.markdown(result)
                
                ### Download button
                st.download_button(
                    label="📥 Download Analysis",
                    data=result,
                    file_name="malware_analysis.txt",
                    mime="text/plain"
                )