frostyripper1
472 lines
12 KiB
Rust
472 lines
12 KiB
Rust
use anyhow::{Context, Result};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::io::Write;
|
|
use std::time::Duration;
|
|
use tokio::io::AsyncReadExt;
|
|
|
|
#[derive(Serialize, Deserialize, Debug, Clone)]
|
|
pub struct ExecutionResult {
|
|
pub command: String,
|
|
pub exit_code: Option<i32>,
|
|
pub stdout: String,
|
|
pub stderr: String,
|
|
pub duration_ms: u64,
|
|
pub success: bool,
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
#[derive(Serialize, Deserialize, Debug, Clone)]
|
|
pub struct ScriptExecution {
|
|
pub language: String,
|
|
pub code: String,
|
|
pub result: ExecutionResult,
|
|
}
|
|
|
|
pub async fn execute_command(
|
|
cmd: &str,
|
|
args: &[&str],
|
|
timeout_secs: u64,
|
|
) -> Result<ExecutionResult> {
|
|
let start = std::time::Instant::now();
|
|
|
|
let mut child = tokio::process::Command::new(cmd)
|
|
.args(args)
|
|
.stdout(std::process::Stdio::piped())
|
|
.stderr(std::process::Stdio::piped())
|
|
.spawn()
|
|
.with_context(|| format!("Failed to spawn: {} {}", cmd, args.join(" ")))?;
|
|
|
|
let timeout = Duration::from_secs(timeout_secs);
|
|
|
|
match tokio::time::timeout(timeout, child.wait()).await {
|
|
Ok(Ok(status)) => {
|
|
let elapsed = start.elapsed();
|
|
let mut stdout_buf = String::new();
|
|
let mut stderr_buf = String::new();
|
|
if let Some(mut out) = child.stdout.take() {
|
|
let _ = out.read_to_string(&mut stdout_buf).await;
|
|
}
|
|
if let Some(mut err) = child.stderr.take() {
|
|
let _ = err.read_to_string(&mut stderr_buf).await;
|
|
}
|
|
Ok(ExecutionResult {
|
|
command: format!("{} {}", cmd, args.join(" ")),
|
|
exit_code: status.code(),
|
|
stdout: stdout_buf.trim().to_string(),
|
|
stderr: stderr_buf.trim().to_string(),
|
|
duration_ms: elapsed.as_millis() as u64,
|
|
success: status.success(),
|
|
})
|
|
}
|
|
Ok(Err(e)) => Err(anyhow::anyhow!("Process failed: {}", e)),
|
|
Err(_) => {
|
|
let _ = child.kill().await;
|
|
let elapsed = start.elapsed();
|
|
Ok(ExecutionResult {
|
|
command: format!("{} {}", cmd, args.join(" ")),
|
|
exit_code: None,
|
|
stdout: String::new(),
|
|
stderr: format!("Killed after {}s timeout", timeout_secs),
|
|
duration_ms: elapsed.as_millis() as u64,
|
|
success: false,
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn detect_language(code: &str) -> &'static str {
|
|
let trimmed = code.trim();
|
|
let first_line = trimmed.lines().next().unwrap_or("").trim();
|
|
|
|
if first_line.starts_with("#!/") {
|
|
if first_line.contains("python") || first_line.contains("python3") {
|
|
return "python";
|
|
}
|
|
if first_line.contains("bash") {
|
|
return "bash";
|
|
}
|
|
if first_line.contains("sh") && !first_line.contains("bash") {
|
|
return "bash";
|
|
}
|
|
if first_line.contains("ruby") {
|
|
return "ruby";
|
|
}
|
|
if first_line.contains("perl") {
|
|
return "perl";
|
|
}
|
|
if first_line.contains("node") {
|
|
return "node";
|
|
}
|
|
}
|
|
|
|
if trimmed.contains("fn main")
|
|
|| (trimmed.contains("println!") && trimmed.contains("fn "))
|
|
|| trimmed.contains("use std::")
|
|
{
|
|
return "rust";
|
|
}
|
|
|
|
if (trimmed.contains("package main") && trimmed.contains("func main"))
|
|
|| trimmed.contains("import (")
|
|
|| trimmed.starts_with("package ")
|
|
{
|
|
return "go";
|
|
}
|
|
|
|
if trimmed.contains("#include <") {
|
|
if trimmed.contains("iostream")
|
|
|| trimmed.contains("std::")
|
|
|| trimmed.contains("using namespace")
|
|
{
|
|
return "cpp";
|
|
}
|
|
return "c";
|
|
}
|
|
|
|
let ps_score = [
|
|
"Write-Host",
|
|
"Get-Process",
|
|
"Get-Service",
|
|
"Get-ChildItem",
|
|
"Param(",
|
|
"Write-Output",
|
|
"Write-Error",
|
|
"Get-Item",
|
|
"Set-Item",
|
|
"New-Item",
|
|
"Remove-Item",
|
|
"ForEach-Object",
|
|
"Where-Object",
|
|
"Select-Object",
|
|
"| ForEach-Object",
|
|
"foreach ($",
|
|
"function ",
|
|
"Export-Csv",
|
|
]
|
|
.iter()
|
|
.filter(|&&kw| trimmed.contains(kw))
|
|
.count();
|
|
if ps_score >= 2 || (trimmed.starts_with("#requires") || trimmed.starts_with("<#")) {
|
|
return "powershell";
|
|
}
|
|
if trimmed.contains("$") && (trimmed.contains("Write-Host") || trimmed.contains("Write-Output"))
|
|
{
|
|
return "powershell";
|
|
}
|
|
|
|
let py_score = [
|
|
"import ",
|
|
"from ",
|
|
"def ",
|
|
"class ",
|
|
"print(",
|
|
"if __name__",
|
|
"elif ",
|
|
"except ",
|
|
"with ",
|
|
"as ",
|
|
"lambda ",
|
|
"yield ",
|
|
"async def",
|
|
"await ",
|
|
]
|
|
.iter()
|
|
.filter(|&&kw| trimmed.contains(kw))
|
|
.count();
|
|
if py_score >= 3 {
|
|
return "python";
|
|
}
|
|
|
|
let js_score = [
|
|
"require(",
|
|
"console.log",
|
|
"module.exports",
|
|
"process.env",
|
|
"=>",
|
|
"const ",
|
|
"let ",
|
|
"var ",
|
|
"function(",
|
|
"async (",
|
|
"document.",
|
|
"window.",
|
|
"export default",
|
|
"import ",
|
|
]
|
|
.iter()
|
|
.filter(|&&kw| trimmed.contains(kw))
|
|
.count();
|
|
if js_score >= 2 {
|
|
return "node";
|
|
}
|
|
|
|
let rb_score = [
|
|
"def ",
|
|
"\nend",
|
|
"puts ",
|
|
"require '",
|
|
"attr_accessor",
|
|
"attr_reader",
|
|
"attr_writer",
|
|
"do |",
|
|
"=>",
|
|
"::",
|
|
]
|
|
.iter()
|
|
.filter(|&&kw| trimmed.contains(kw))
|
|
.count();
|
|
if rb_score >= 2 {
|
|
return "ruby";
|
|
}
|
|
|
|
let pl_score = [
|
|
"use strict",
|
|
"use warnings",
|
|
"my $",
|
|
"sub ",
|
|
"shift",
|
|
"print \"",
|
|
"print '",
|
|
"chomp",
|
|
"foreach my",
|
|
"while (<",
|
|
]
|
|
.iter()
|
|
.filter(|&&kw| trimmed.contains(kw))
|
|
.count();
|
|
if pl_score >= 2 {
|
|
return "perl";
|
|
}
|
|
|
|
let sh_score = [
|
|
"if [",
|
|
"if [[",
|
|
"$#",
|
|
"\nfi",
|
|
"\ndone",
|
|
"\nthen",
|
|
"case ",
|
|
"esac",
|
|
">&2",
|
|
"2>&1",
|
|
"#!/",
|
|
"echo \"",
|
|
"while read",
|
|
"for i in",
|
|
"do\n",
|
|
"$((",
|
|
"${",
|
|
]
|
|
.iter()
|
|
.filter(|&&kw| trimmed.contains(kw))
|
|
.count();
|
|
if sh_score >= 2 {
|
|
return "bash";
|
|
}
|
|
|
|
let cmd_score = [
|
|
"@echo",
|
|
"echo off",
|
|
"setlocal",
|
|
"endlocal",
|
|
"%",
|
|
"REM ",
|
|
"rem ",
|
|
"color ",
|
|
"title ",
|
|
"exit /b",
|
|
"if not",
|
|
"if defined",
|
|
"set \"",
|
|
">nul",
|
|
]
|
|
.iter()
|
|
.filter(|&&kw| trimmed.contains(kw))
|
|
.count();
|
|
if cmd_score >= 2 {
|
|
return "cmd";
|
|
}
|
|
|
|
"python"
|
|
}
|
|
|
|
pub async fn execute_script(
|
|
language: &str,
|
|
code: &str,
|
|
timeout_secs: u64,
|
|
) -> Result<ExecutionResult> {
|
|
let (interpreter, ext) = match language.to_lowercase().as_str() {
|
|
"python" | "py" => ("python", ".py"),
|
|
"bash" | "sh" => ("bash", ".sh"),
|
|
"powershell" | "ps1" => ("powershell", ".ps1"),
|
|
"cmd" | "bat" => ("cmd", ".bat"),
|
|
"ruby" => ("ruby", ".rb"),
|
|
"perl" => ("perl", ".pl"),
|
|
"node" | "javascript" | "js" => ("node", ".js"),
|
|
"rust" => return compile_and_run_rust(code, timeout_secs).await,
|
|
"go" | "golang" => return compile_and_run_go(code, timeout_secs).await,
|
|
"c" => return compile_and_run_c_cpp(code, timeout_secs, "c").await,
|
|
"cpp" | "c++" | "cxx" => return compile_and_run_c_cpp(code, timeout_secs, "cpp").await,
|
|
_ => return Err(anyhow::anyhow!("Unsupported language: {}. Supported: python, bash, powershell, cmd, ruby, perl, node, rust, go, c, cpp", language)),
|
|
};
|
|
|
|
let temp_dir = std::env::temp_dir().join("airust_scripts");
|
|
std::fs::create_dir_all(&temp_dir)?;
|
|
let script_path = temp_dir.join(format!("script_{}_{}", std::process::id(), ext));
|
|
|
|
let mut file = std::fs::File::create(&script_path)?;
|
|
file.write_all(code.as_bytes())?;
|
|
drop(file);
|
|
|
|
let result = if interpreter == "cmd" {
|
|
execute_command("cmd", &["/C", script_path.to_str().unwrap()], timeout_secs).await
|
|
} else {
|
|
execute_command(interpreter, &[script_path.to_str().unwrap()], timeout_secs).await
|
|
};
|
|
|
|
let _ = std::fs::remove_file(&script_path);
|
|
result
|
|
}
|
|
|
|
async fn compile_and_run_rust(code: &str, timeout_secs: u64) -> Result<ExecutionResult> {
|
|
let temp_dir = std::env::temp_dir().join("airust_scripts");
|
|
std::fs::create_dir_all(&temp_dir)?;
|
|
|
|
let source_path = temp_dir.join(format!("script_{}.rs", std::process::id()));
|
|
let bin_path = temp_dir.join(format!("script_{}", std::process::id()));
|
|
|
|
let mut file = std::fs::File::create(&source_path)?;
|
|
file.write_all(code.as_bytes())?;
|
|
drop(file);
|
|
|
|
let compile = execute_command(
|
|
"rustc",
|
|
&[
|
|
source_path.to_str().unwrap(),
|
|
"-o",
|
|
bin_path.to_str().unwrap(),
|
|
"--quiet",
|
|
],
|
|
timeout_secs,
|
|
)
|
|
.await?;
|
|
if !compile.success {
|
|
return Ok(ExecutionResult {
|
|
command: "rustc (compile)".to_string(),
|
|
exit_code: compile.exit_code,
|
|
stdout: String::new(),
|
|
stderr: compile.stderr,
|
|
duration_ms: compile.duration_ms,
|
|
success: false,
|
|
});
|
|
}
|
|
|
|
let run_result = execute_command(bin_path.to_str().unwrap(), &[], timeout_secs).await;
|
|
|
|
let _ = std::fs::remove_file(&source_path);
|
|
let _ = std::fs::remove_file(&bin_path);
|
|
run_result
|
|
}
|
|
|
|
async fn compile_and_run_go(code: &str, timeout_secs: u64) -> Result<ExecutionResult> {
|
|
let temp_dir = std::env::temp_dir().join("airust_scripts");
|
|
std::fs::create_dir_all(&temp_dir)?;
|
|
|
|
let source_path = temp_dir.join(format!("script_{}.go", std::process::id()));
|
|
let bin_path = temp_dir.join(format!("script_{}.exe", std::process::id()));
|
|
|
|
let mut file = std::fs::File::create(&source_path)?;
|
|
file.write_all(code.as_bytes())?;
|
|
drop(file);
|
|
|
|
let compile = execute_command(
|
|
"go",
|
|
&[
|
|
"build",
|
|
"-o",
|
|
bin_path.to_str().unwrap(),
|
|
source_path.to_str().unwrap(),
|
|
],
|
|
timeout_secs,
|
|
)
|
|
.await?;
|
|
if !compile.success {
|
|
return Ok(ExecutionResult {
|
|
command: "go build".to_string(),
|
|
exit_code: compile.exit_code,
|
|
stdout: String::new(),
|
|
stderr: compile.stderr,
|
|
duration_ms: compile.duration_ms,
|
|
success: false,
|
|
});
|
|
}
|
|
|
|
let run_result = execute_command(bin_path.to_str().unwrap(), &[], timeout_secs).await;
|
|
|
|
let _ = std::fs::remove_file(&source_path);
|
|
let _ = std::fs::remove_file(&bin_path);
|
|
run_result
|
|
}
|
|
|
|
async fn compile_and_run_c_cpp(
|
|
code: &str,
|
|
timeout_secs: u64,
|
|
lang: &str,
|
|
) -> Result<ExecutionResult> {
|
|
let temp_dir = std::env::temp_dir().join("airust_scripts");
|
|
std::fs::create_dir_all(&temp_dir)?;
|
|
|
|
let ext = if lang == "cpp" { ".cpp" } else { ".c" };
|
|
let source_path = temp_dir.join(format!("script_{}{}", std::process::id(), ext));
|
|
let bin_path = temp_dir.join(format!("script_{}.exe", std::process::id()));
|
|
|
|
let mut file = std::fs::File::create(&source_path)?;
|
|
file.write_all(code.as_bytes())?;
|
|
drop(file);
|
|
|
|
let compiler = if lang == "cpp" { "g++" } else { "gcc" };
|
|
let compile = execute_command(
|
|
compiler,
|
|
&[
|
|
source_path.to_str().unwrap(),
|
|
"-o",
|
|
bin_path.to_str().unwrap(),
|
|
],
|
|
timeout_secs,
|
|
)
|
|
.await?;
|
|
if !compile.success {
|
|
return Ok(ExecutionResult {
|
|
command: format!("{} (compile)", compiler),
|
|
exit_code: compile.exit_code,
|
|
stdout: String::new(),
|
|
stderr: compile.stderr,
|
|
duration_ms: compile.duration_ms,
|
|
success: false,
|
|
});
|
|
}
|
|
|
|
let run_result = execute_command(bin_path.to_str().unwrap(), &[], timeout_secs).await;
|
|
|
|
let _ = std::fs::remove_file(&source_path);
|
|
let _ = std::fs::remove_file(&bin_path);
|
|
run_result
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
pub fn list_supported_languages() -> Vec<String> {
|
|
vec![
|
|
"python".to_string(),
|
|
"bash".to_string(),
|
|
"powershell".to_string(),
|
|
"cmd".to_string(),
|
|
"node".to_string(),
|
|
"ruby".to_string(),
|
|
"perl".to_string(),
|
|
"rust".to_string(),
|
|
"go".to_string(),
|
|
"c".to_string(),
|
|
"cpp".to_string(),
|
|
]
|
|
}
|