java_decompiler_ollama/
lib.rs

1use std::env;
2use std::process::Command;
3use reqwest::blocking::Client;
4use reqwest::header::{HeaderMap, CONTENT_TYPE};
5use serde_json::{json, Value};
6use thiserror::Error;
7
8/// Errors that can occur during the translation process.
9#[derive(Debug, Error)]
10pub enum TranslationError {
11    /// Error reading the file.
12    #[error("Error reading the file: {0}")]
13    FileReadError(#[from] std::io::Error),
14    /// Error sending request to Ollama or retrieving response.
15    #[error("Error sending request to Ollama or retrieving response: {0}")]
16    RequestError(#[from] reqwest::Error),
17    /// Ollama error with HTTP status code.
18    #[error("Ollama error: HTTP status {0}")]
19    OllamaError(u16),
20    /// Error parsing JSON.
21    #[error("Error parsing JSON: {0}")]
22    JsonError(#[from] serde_json::Error),
23    /// Error running javap.
24    #[error("Error running javap: {0}")]
25    JavapError(String),
26}
27
28/// Configuration struct for the translation process.
29struct Config {
30    model: String,
31    url: String,
32}
33
34impl Config {
35    /// Creates a new Config from environment variables.
36    fn from_env() -> Self {
37        Self {
38            model: env::var("OLLAMA_MODEL").unwrap_or_else(|_| "qwen2.5-coder".to_string()),
39            url: env::var("OLLAMA_URL").unwrap_or_else(|_| "http://localhost:11434/api/generate".to_string()),
40        }
41    }
42}
43
44/// Creates the headers required for the HTTP request.
45fn create_headers() -> HeaderMap {
46    let mut headers = HeaderMap::new();
47    headers.insert(CONTENT_TYPE, "application/json".parse().unwrap());
48    headers
49}
50
51/// Parses the response from the Ollama API and extracts the translated code.
52fn parse_response(response_text: &str) -> Result<String, TranslationError> {
53    let mut result = String::new();
54    let mut skip_first_line = true;
55    for line in response_text.lines() {
56        if let Ok(json_value) = serde_json::from_str::<Value>(line) {
57            if let Some(response_content) = json_value.get("response") {
58                if let Some(content) = response_content.as_str() {
59                    if content.contains("<|im_start|>") || content.contains("```") {
60                        continue;
61                    }
62                    if skip_first_line && !content.trim().is_empty() {
63                        skip_first_line = false;
64                        continue;
65                    }
66                    result.push_str(content);
67                }
68            }
69        }
70    }
71
72    let clean_result = result
73        .lines()
74        .skip_while(|line| line.trim().is_empty() || line.contains("<|im_start|>") || line.contains("```"))
75        .take_while(|line| !line.contains("```"))
76        .collect::<Vec<&str>>()
77        .join("\n");
78
79    Ok(clean_result.trim().to_string())
80}
81
82/// Disassembles a Java class file using javap and translates the bytecode to Java source code using the Ollama API.
83///
84/// # Arguments
85///
86/// * `class_file` - A string slice that holds the path to the Java class file to be disassembled.
87///
88/// # Returns
89///
90/// * `Ok(String)` - The translated Java source code as a string.
91/// * `Err(TranslationError)` - An error that occurred during the translation process.
92pub fn translate_java_class(class_file: &str) -> Result<String, TranslationError> {
93    // Appelle javap pour désassembler la classe Java
94    let output = Command::new("javap")
95        .arg("-c")
96        .arg(class_file)
97        .output()
98        .map_err(|e| TranslationError::JavapError(e.to_string()))?;
99
100    if !output.status.success() {
101        return Err(TranslationError::JavapError(String::from_utf8_lossy(&output.stderr).to_string()));
102    }
103
104    let javap_output = String::from_utf8_lossy(&output.stdout).to_string();
105
106    // Envoie le bytecode désassemblé à Ollama
107    let config = Config::from_env();
108    let client = Client::new();
109    let body = json!({
110        "model": config.model,
111        "prompt": format!(
112            "Translate the following bytecode to Java source code. Output only the code:\n{}",
113            javap_output
114        ),
115        "stream": true
116    });
117
118    let headers = create_headers();
119    let response = client
120        .post(&config.url)
121        .headers(headers)
122        .json(&body)
123        .send()?;
124
125    if !response.status().is_success() {
126        return Err(TranslationError::OllamaError(response.status().as_u16()));
127    }
128
129    let response_text = response.text()?;
130    let translated_code = parse_response(&response_text)?;
131
132    Ok(translated_code)
133}
134