UI.py

import streamlit as st
import requests, base64
from google import genai
from secret_config import GEMINI_API_KEY

# Set up the Gemini API key and client
Gemini_api_key = GEMINI_API_KEY
client = genai.Client(api_key=Gemini_api_key)

def get_github_file_content(owner, repo, file_path, branch="main"):
    """
    Fetch the content of a file from a GitHub repository.
    """
    url = f"https://api.github.com/repos/{owner}/{repo}/contents/{file_path}?ref={branch}"
    headers = {"User-Agent": "Mozilla/5.0"}
    response = requests.get(url, headers=headers)
    if response.status_code == 200:
        file_data = response.json()
        content = base64.b64decode(file_data["content"]).decode("utf-8")
        return content
    else:
        return f"Error: {response.status_code}, {response.text}"

def analyze_code(code):
    """
    Use Gemini to analyze the provided code and output vulnerabilities in a concise table.
    """
    context = (
        "You are a security expert specializing in static code analysis. Your task is to analyze the following code and identify potential vulnerabilities. "
        "For each vulnerability you detect, provide a concise summary in a table format with exactly three columns: 'Vulnerability', 'Location' (e.g., line number or code section), "
        "and 'Risk Level' (e.g., Low, Medium, or High). Do not include detailed explanations or descriptions—only list the vulnerability name, its location, and the risk level.\n"
        "Code to analyze:\n"
    )
    full_prompt = context + code
    response = client.models.generate_content(
        model="gemini-1.5-flash",
        contents=full_prompt,
    )
    return response.text

def main():
    st.title("Gemini Vulnerability Analyzer")
    st.write("Analyze your code for potential vulnerabilities using Gemini AI.")

    # Sidebar option to select the analysis mode
    mode = st.sidebar.radio("Choose Analysis Mode", ["Analyze Code", "Analyze GitHub Repository"])
    
    if mode == "Analyze Code":
        code = st.text_area("Enter your code here:", height=300)
        if st.button("Analyze"):
            if code.strip():
                with st.spinner("Analyzing code..."):
                    analysis = analyze_code(code)
                st.subheader("Vulnerability Analysis:")
                st.code(analysis)
            else:
                st.error("Please paste some code to analyze.")
    
    elif mode == "Analyze GitHub Repository":
        st.subheader("GitHub Repository Settings")
        owner = st.text_input("Repository Owner", "KshshVrma")
        repo = st.text_input("Repository Name", "Student-Mangement-System")
        branch = st.text_input("Branch", "main")
        
        st.write("Fetching file list from the repository...")
        tree_url = f"https://api.github.com/repos/{owner}/{repo}/git/trees/{branch}?recursive=1"
        headers = {"User-Agent": "Mozilla/5.0"}
        response = requests.get(tree_url, headers=headers)
        
        if response.status_code == 200:
            data = response.json()
            files = [item["path"] for item in data["tree"] if item["type"] == "blob"]
            if files:
                selected_file = st.selectbox("Select a file to analyze", files)
                if st.button("Analyze Selected File"):
                    code_content = get_github_file_content(owner, repo, selected_file, branch)
                    st.subheader("File Content:")
                    st.code(code_content)
                    with st.spinner("Analyzing file..."):
                        analysis = analyze_code(code_content)
                    st.subheader("Vulnerability Analysis:")
                    st.code(analysis)
            else:
                st.error("No files found in the repository.")
        else:
            st.error(f"Error fetching repository: {response.status_code} - {response.text}")

if __name__ == "__main__":
    main()