Spaces:
Running
Running
File size: 4,450 Bytes
3933ced 2eb374f 3933ced 2eb374f 3933ced | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | 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)} |