From 7fe74629da46731386c3315b30601cadca27cb76 Mon Sep 17 00:00:00 2001
From: toast <toast@toast-server.net>
Date: Tue, 5 Dec 2023 14:18:30 +1100
Subject: [PATCH] Add `/eval` command

---
 .vscode/settings.json |  3 ++-
 Cargo.lock            |  1 +
 Cargo.toml            |  1 +
 src/commands/eval.rs  | 56 +++++++++++++++++++++++++++++++++++++++++++
 src/commands/mod.rs   |  1 +
 src/main.rs           |  3 ++-
 6 files changed, 63 insertions(+), 2 deletions(-)
 create mode 100644 src/commands/eval.rs

diff --git a/.vscode/settings.json b/.vscode/settings.json
index 7e94c91..2938df3 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -1,5 +1,6 @@
 {
   "rust-analyzer.linkedProjects": [
     "./Cargo.toml"
-  ]
+  ],
+  "rust-analyzer.showUnlinkedFileNotification": false
 }
diff --git a/Cargo.lock b/Cargo.lock
index 82c2cea..b3f0caa 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1146,6 +1146,7 @@ version = "0.1.0"
 dependencies = [
  "poise",
  "serenity 0.12.0",
+ "tempfile",
  "tokio",
  "tracing",
  "tracing-subscriber",
diff --git a/Cargo.toml b/Cargo.toml
index 389f3d7..4f74436 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -8,6 +8,7 @@ edition = "2021"
 [dependencies]
 poise = "0.5.7"
 serenity = "0.12.0"
+tempfile = "3.8.1"
 tokio = { version = "1.34.0", features = ["macros", "signal", "rt-multi-thread"] }
 tracing = "0.1.40"
 tracing-subscriber = "0.3.18"
diff --git a/src/commands/eval.rs b/src/commands/eval.rs
new file mode 100644
index 0000000..375240b
--- /dev/null
+++ b/src/commands/eval.rs
@@ -0,0 +1,56 @@
+use crate::Error;
+
+use poise::serenity_prelude::UserId;
+use std::os::unix::fs::PermissionsExt;
+use std::process::Command;
+use std::io::Write;
+
+const WHITELISTED_USERS: &[UserId] = &[
+  poise::serenity_prelude::UserId(190407856527376384)
+];
+
+/// Evaluate a piece of code
+#[poise::command(slash_command)]
+pub async fn eval(
+  ctx: poise::Context<'_, (), Error>,
+  #[description = "The Rust code to evaluate"] code: String
+) -> Result<(), Error> {
+  if !WHITELISTED_USERS.contains(&ctx.author().id) {
+    ctx.reply("Whitelisted users can only use this command!").await?;
+    return Ok(());
+  }
+
+  // Create a temp directory
+  let dir = tempfile::tempdir()?;
+  let file_path = dir.path().join("temp.rs");
+
+  {
+    let mut file = std::fs::File::create(&file_path)?;
+    writeln!(file, "{}", code)?;
+  }
+
+  // Compile
+  let compiled_path = dir.path().join("temp");
+  let output = Command::new("rustc").arg(&file_path).arg("-o").arg(&compiled_path).output()?;
+
+  if !output.status.success() {
+    ctx.reply(format!("Compilation failed:\n```{}```", String::from_utf8_lossy(&output.stderr))).await?;
+    return Ok(());
+  }
+
+  // Update binary's permissions before execution stage
+  let permissions = std::fs::Permissions::from_mode(0o755);
+  let compiled_path = dir.path().join("temp");
+  std::fs::set_permissions(&compiled_path, permissions)?;
+
+  // If success, run it.
+  let output = Command::new(compiled_path).output()?;
+
+  if !output.status.success() {
+    ctx.reply(format!("Execution failed:\n```{}```", String::from_utf8_lossy(&output.stderr))).await?;
+    return Ok(());
+  }
+
+  ctx.reply(format!("Code output:\n```rs\n{}```", String::from_utf8_lossy(&output.stdout))).await?;
+  Ok(())
+}
diff --git a/src/commands/mod.rs b/src/commands/mod.rs
index a766209..85dd8f6 100644
--- a/src/commands/mod.rs
+++ b/src/commands/mod.rs
@@ -1 +1,2 @@
 pub mod ping;
+pub mod eval;
diff --git a/src/main.rs b/src/main.rs
index 6694ced..b8b68da 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -37,7 +37,8 @@ async fn main() {
     .intents(serenity::GatewayIntents::MESSAGE_CONTENT | serenity::GatewayIntents::GUILDS)
     .options(poise::FrameworkOptions {
       commands: vec![
-        commands::ping::ping()
+        commands::ping::ping(),
+        commands::eval::eval()
       ],
       pre_command: |ctx| {
         Box::pin(async move {