style_engine.py

import os
import logging
from google import genai
from google.genai import types
import config

# Configure Logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class StyleEngine:
    def __init__(self):
        self.api_key = os.getenv("GEMINI_API_KEY")
        self.client = None

        if self.api_key:
            try:
                # Initialize Client (New SDK)
                self.client = genai.Client(api_key=self.api_key)
                logger.info("✅ Gemini Client Initialized Successfully.")
            except Exception as e:
                logger.error(f"❌ Failed to initialize Gemini Client: {e}")
        else:
            logger.warning("⚠️ GEMINI_API_KEY not found. Style transfer will be disabled.")

    def juliet_style_transfer(self, source_code):
        """
        Rewrites code to match the NIST Juliet Test Suite style.
        """
        if not self.client:
            logger.error("❌ Attempted to call Gemini but client is not initialized.")
            return {
                "status": "error",
                "message": "Gemini API Key is missing."
            }

        logger.info("🚀 Preparing to call Gemini API for Style Transfer...")

        prompt = f"""
        You are an expert C Code Normalizer specialized in the NIST Juliet Test Suite standards.
        Your goal is to rewrite arbitrary C code to strictly mimic the style, structure, and variable naming of a NIST Juliet CWE test case.

        ### STYLE GUIDE:
        1.  **Variable Names:** Use 'data', 'dataBuffer', 'dataBadBuffer', 'dataGoodBuffer' etc to match the Juliet Style.
        2.  **Function Names:** Rename the main logic function to 'CWE<ID>_<Name>__bad()'to match the Juliet Style.
        3.  **Comments:** Use Juliet-style headers like '/* POTENTIAL FLAW: ... */'.
        4.  **Headers:** Use standard headers (<stdio.h>, <stdlib.h>, <string.h>). DO NOT include 'std_testcase.h'.
        5.  **Formatting:** 4-space indentation. Braces on new lines (Allman style).

        ### CRITICAL RULES:
        1.  **PRESERVE VULNERABILITIES:** Do not fix bugs. If the logic causes a crash, keep it.
        2.  **PRESERVE LOGIC:** The control flow must remain exactly the same.
        3.  **STANDALONE:** The code must compile with 'gcc source.c'.

        ### EXAMPLES:

        --- Example 1 ---
        Input:
        #include <stdio.h>
        #include <string.h>
        #include <stdlib.h>

        void func(char *str) {{
        char buf[10];
        strcpy(buf, str);
        }}
        
        int main() {{
        func("loooooooooongstring");
        }}

        Output:
        #include <stdio.h>
        #include <string.h>
        #include <stdlib.h>

        void CWE121_Stack_Based_Buffer_Overflow__bad()
        {{
            char * data;
            char dataBuffer[10];
            /* POTENTIAL FLAW: Input string might be larger than buffer */
            data = "loooooooooongstring"; 
            strcpy(dataBuffer, data);
        }}

        int main(int argc, char * argv[])
        {{
            CWE121_Stack_Based_Buffer_Overflow__bad();
            return 0;
        }}
        ------------------

        ### YOUR TASK:
        Convert the following code to Juliet Style. 
        RETURN ONLY THE RAW C CODE. NO MARKDOWN (```c), NO EXPLANATIONS.

        Input Code:
        {source_code}
        """

        try:
            logger.info("📡 Sending request to Gemini...")
            
            # CALL API (New SDK Syntax)
            response = self.client.models.generate_content(
                model="gemini-3-flash-preview", 
                contents=prompt
            )
            
            logger.info("✅ Gemini API Call Successful. Processing response...")

            # Cleanup output
            if response.text:
                clean_code = response.text.replace("```c", "").replace("```cpp", "").replace("```", "").strip()
                logger.info("✨ Style Transfer Complete.")
                return {
                    "status": "success",
                    "juliet_code": clean_code
                }
            
            logger.warning("⚠️ Gemini returned an empty response.")
            return {"status": "error", "message": "Empty response from Gemini"}

        except Exception as e:
            logger.error(f"❌ Gemini API Error: {e}")
            return {"status": "error", "message": str(e)}