From 0921786cc1bdbd246775099d791c158dbefd1126 Mon Sep 17 00:00:00 2001 From: mii Date: Mon, 24 Oct 2022 14:22:14 +0000 Subject: [PATCH] init --- .gitignore | 4 + Cargo.toml | 33 +++++++ sample-config.yaml | 180 +++++++++++++++++++++++++++++++++++ src/config.rs | 22 +++++ src/docker.rs | 222 +++++++++++++++++++++++++++++++++++++++++++ src/event_handler.rs | 162 +++++++++++++++++++++++++++++++ src/language.rs | 46 +++++++++ src/main.rs | 80 ++++++++++++++++ 8 files changed, 749 insertions(+) create mode 100644 .gitignore create mode 100644 Cargo.toml create mode 100644 sample-config.yaml create mode 100644 src/config.rs create mode 100644 src/docker.rs create mode 100644 src/event_handler.rs create mode 100644 src/language.rs create mode 100644 src/main.rs diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..4140850 --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +/target +/code +config.yaml +Cargo.lock \ No newline at end of file diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..362f4d0 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,33 @@ +[package] +name = "rs-docker-bot" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +bollard = "0.13" +futures-util = "0.3" +serde_json = "1.0" +serde_yaml = "0.9" +serde = { version = "1.0", features = ["derive"] } +uuid = { version = "0.8", features = ["serde", "v4"] } +regex = "1" +phf = { version = "0.11.1", features = ["macros"] } +flate2 = "1" +tar = "0.4" +memfile = "0.2" +strum = "0.24" +strum_macros = "0.24" +log = "0.4.0" +env_logger = "0.9.0" + +[dependencies.serenity] +default-features = true +features = ["builder", "cache", "client", "gateway", "model", "utils", "unstable_discord_api", "collector", "rustls_backend", "framework"] +git = "https://github.com/serenity-rs/serenity" +branch = "next" + +[dependencies.tokio] +version = "1.0" +features = ["macros", "rt-multi-thread"] diff --git a/sample-config.yaml b/sample-config.yaml new file mode 100644 index 0000000..a72f517 --- /dev/null +++ b/sample-config.yaml @@ -0,0 +1,180 @@ +token: TOKEN +prefix: PREFIX +languages: +- name: Ruby + code: + - Ruby + - ruby + - rb + extension: rb + path: "{file}" + run_command: "ruby ./{file}" + image: ruby + +- name: Python + code: + - Python + - python + - py + extension: py + path: "{file}" + run_command: "python ./{file}" + image: "python:3" + +- name: JavaScript + code: + - JavaScript + - javascript + - js + extension: js + path: "{file}" + run_command: "node ./{file}" + image: node + +- name: "C++" + code: + - cpp + - c++ + - c + extension: cpp + path: "{file}" + compile_command: "gcc {file} -o program" + run_command: "./program" + image: gcc + +- name: Java + code: + - java + extension: java + path: "Main.java" + compile_command: "javac Main.java" + run_command: "java Main" + image: "openjdk:7" + +- name: Kotlin + code: + - Kotlin + - kt + extension: kt + path: "Main.kt" + compile_command: "kotlinc Main.kt" + run_command: "kotlin MainKt" + image: "zenika/kotlin:latest" + +- name: Julia + code: + - julia + - jl + extension: jl + path: "{file}" + run_command: "julia ./{file}" + image: julia + +- name: Rust + code: + - rust + - rs + extension: rs + path: "{file}" + compile_command: "rustc {file} -o program" + run_command: "./program" + image: rust + +- name: PHP + code: + - php + extension: php + path: "{file}" + run_command: "php ./{file}" + image: "php:7.4-cli" + +- name: Go + code: + - go + extension: go + path: "{file}" + run_command: "go run ./{file}" + image: golang + +- name: GP + code: + - gp + - pari + extension: gp + path: "{file}" + run_command: "gp -q ./{file}" + image: "pascalmolin/parigp-full" + +- name: Bash + code: + - bash + - sh + extension: sh + path: "{file}" + run_command: "bash ./{file}" + image: ubuntu + +- name: Maxima + code: + - maxima + - mc + - mac + extension: mac + path: "{file}" + run_command: "maxima --very-quiet ./{file}" + image: "jgoldfar/maxima-docker:debian-latest" + +- name: ăȘでしこ + code: + - nadesiko + - nako + extension: nako3 + path: "{file}" + run_command: "nadesiko /{file}" + image: "esolang/nadesiko" + +- name: Fortran + code: + - fortran + extension: f + path: "{file}" + compile_command: "gfortran -o program /{file}" + run_command: "./program" + image: "nacyot/fortran-gfortran:apt" + +- name: R + code: + - R + - r + extension: R + path: "{file}" + run_command: "Rscript {file}" + image: "r-base" + +- name: "x86 ASM NASM" + code: + - x86 + - asm + - x86asm + extension: "x86.asm" + path: "{file}" + run_command: "x86asm-nasm /{file}" + image: "esolang/x86asm-nasm" + +- name: Elixir + code: + - elixir + - exs + extension: exs + path: "{file}" + run_command: "elixir /{file}" + image: "esolang/elixir" + +- name: Haskell + code: + - haskell + - hs + extension: hs + path: "{file}" + run_command: "haskell /{file}" + image: "esolang/haskell" \ No newline at end of file diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..41e4da5 --- /dev/null +++ b/src/config.rs @@ -0,0 +1,22 @@ +use serde::{Serialize, Deserialize}; + +use crate::language::Language; + +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct Config { + pub token: String, + pub prefix: String, + pub languages: Vec +} + +impl Config { + pub fn get_language(&self, name: &String) -> Option { + for language in self.languages.iter() { + if language.code.contains(name) { + return Some(language.clone()); + } + } + + None + } +} \ No newline at end of file diff --git a/src/docker.rs b/src/docker.rs new file mode 100644 index 0000000..5462d92 --- /dev/null +++ b/src/docker.rs @@ -0,0 +1,222 @@ +use std::{ + fs::File, + io::{Read, Write}, + sync::mpsc::{self, Receiver, Sender}, + time::Duration, +}; + +use bollard::{ + container::{ + CreateContainerOptions, ListContainersOptions, LogOutput, RemoveContainerOptions, + UploadToContainerOptions, + }, + exec::{CreateExecOptions, StartExecResults}, + service::ContainerSummary, + Docker, +}; +use flate2::{write::GzEncoder, Compression}; +use futures_util::StreamExt; +use tokio::task::JoinHandle; + +use crate::language::Language; + +pub async fn docker_ps() -> Vec { + let docker = Docker::connect_with_local_defaults().unwrap(); + + let options = ListContainersOptions:: { + ..Default::default() + }; + + let list = docker.list_containers(Some(options)).await; + + list.unwrap() +} + +pub struct Container { + pub id: String, + pub name: String, + pub language: Option, +} + +impl Container { + pub async fn from_language(language: Language) -> Self { + let docker = Docker::connect_with_local_defaults().unwrap(); + let config = language.get_container_option(); + let name = format!("dockerbot-{}", uuid::Uuid::new_v4().to_string()); + + let id = docker + .create_container(Some(CreateContainerOptions { name: name.clone() }), config) + .await + .unwrap() + .id; + + Self { + id, + name, + language: Some(language), + } + } + + pub async fn stop(&self) { + let docker = Docker::connect_with_local_defaults().unwrap(); + + docker + .remove_container( + &self.id, + Some(RemoveContainerOptions { + force: true, + ..Default::default() + }), + ) + .await + .unwrap(); + } + + pub async fn run_code(&self) -> (JoinHandle<()>, Receiver>, Sender<()>) { + let docker = Docker::connect_with_local_defaults().unwrap(); + let language = self.language.clone().unwrap(); + let file_name = format!("{}.{}", self.name, language.extension); + + let exec = docker + .create_exec( + &self.id, + CreateExecOptions { + attach_stdout: Some(true), + attach_stdin: Some(true), + attach_stderr: Some(true), + cmd: Some( + language + .get_run_command(file_name) + .split(" ") + .collect(), + ), + ..Default::default() + }, + ) + .await + .unwrap() + .id; + + let (tx, rx) = mpsc::channel(); + let (end_tx, end_rx) = mpsc::channel::<()>(); + + let handle = tokio::spawn(async move { + if let StartExecResults::Attached { mut output, .. } = + docker.start_exec(&exec, None).await.unwrap() + { + let mut end_flag = false; + while !end_flag { + if let Ok(_) = end_rx.recv_timeout(Duration::from_millis(10)) { + break; + } + if let Ok(res) = + tokio::time::timeout(Duration::from_millis(100), output.next()).await + { + if let Some(Ok(msg)) = res { + tx.send(Some(msg)).unwrap(); + } else { + end_flag = true; + } + } + } + } else { + unreachable!(); + } + + tx.send(None).unwrap(); + }); + + (handle, rx, end_tx) + } + + pub async fn compile(&self) -> Option<(JoinHandle<()>, Receiver>)> { + let docker = Docker::connect_with_local_defaults().unwrap(); + let language = self.language.clone().unwrap(); + let file_name = format!("{}.{}", self.name, language.extension); + + if let Some(compile) = language.get_compile_command(file_name.clone()) { + let exec = docker + .create_exec( + &self.id, + CreateExecOptions { + attach_stdout: Some(true), + attach_stdin: Some(true), + attach_stderr: Some(true), + cmd: Some(compile.split(" ").collect()), + ..Default::default() + }, + ) + .await + .unwrap() + .id; + + let (tx, rx) = mpsc::channel(); + + let handle = tokio::spawn(async move { + if let StartExecResults::Attached { mut output, .. } = + docker.start_exec(&exec, None).await.unwrap() + { + while let Some(Ok(msg)) = output.next().await { + tx.send(Some(msg)).unwrap(); + } + } else { + unreachable!(); + } + + tx.send(None).unwrap(); + }); + + return Some((handle, rx)); + } + return None; + } + + pub async fn upload_file(&self, content: &str, file_name: String) { + let docker = Docker::connect_with_local_defaults().unwrap(); + let path = self.language.clone().unwrap().get_path(file_name.clone()); + + { + let mut f = File::create(&format!("code/{}", path)).unwrap(); + f.write_all(content.as_bytes()).unwrap(); + f.flush().unwrap(); + f.sync_all().unwrap(); + } + + { + let mut f = File::open(&format!("code/{}", path)).unwrap(); + + let tar_gz = File::create(&format!("code/{}.tar.gz", file_name)).unwrap(); + let encoder = GzEncoder::new(tar_gz, Compression::default()); + let mut tar = tar::Builder::new(encoder); + tar.append_file( + path.clone(), + &mut f, + ) + .unwrap(); + tar.finish().unwrap(); + } + + { + let mut f = File::open(&format!("code/{}.tar.gz", file_name)).unwrap(); + let mut contents = Vec::new(); + f.read_to_end(&mut contents).unwrap(); + + docker + .start_container::(&self.id, None) + .await + .unwrap(); + + docker + .upload_to_container( + &self.id, + Some(UploadToContainerOptions { + path: "/", + ..Default::default() + }), + contents.into(), + ) + .await + .unwrap(); + } + } +} diff --git a/src/event_handler.rs b/src/event_handler.rs new file mode 100644 index 0000000..8b51507 --- /dev/null +++ b/src/event_handler.rs @@ -0,0 +1,162 @@ +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +use log::info; + +use regex::Regex; +use serenity::{ + async_trait, + builder::{ + CreateApplicationCommand, CreateApplicationCommandOption, CreateInteractionResponse, + CreateInteractionResponseMessage, EditMessage, + }, + model::prelude::{ + command::{CommandOptionType, CommandType, Command}, + Interaction, Message, Ready, + }, + prelude::{Context, EventHandler}, +}; +use tokio::time::{sleep_until, Instant}; + +use crate::{ + docker::{docker_ps, Container}, + ConfigStorage, +}; + +pub struct Handler; + +#[async_trait] +impl EventHandler for Handler { + async fn message(&self, context: Context, message: Message) { + let regex = Regex::new("^```(?P[0-9a-zA-Z]*)\n(?P(\n|.)+)\n```$").unwrap(); + + let capture = regex.captures(&message.content); + + if let Some(captures) = capture { + let language = captures.name("language").unwrap().as_str(); + let code = captures.name("code").unwrap().as_str(); + + let config = { + let data_read = context.data.read().await; + data_read.get::().expect("Cannot get ConfigStorage.").clone() + }; + + let language = { + let config = config.lock().unwrap(); + config.get_language(&String::from(language)) + }; + + if let Some(language) = language { + let mut message = message + .reply(&context.http, format!("Creating {:?} container.", language.name)) + .await + .unwrap(); + + let container = Container::from_language(language.clone()).await; + let file_name = format!("{}.{}", container.name, language.extension.clone()); + + message + .edit( + &context.http, + EditMessage::new().content(format!("Created: {}", container.id)), + ) + .await + .unwrap(); + + container.upload_file(code, file_name.clone()).await; + + if let Some((compile_handle, compile_rx)) = container.compile().await { + let rx_handle = tokio::spawn(async move { + while let Ok(Some(msg)) = compile_rx.recv() { + print!("{}", msg); + } + }); + + let (_, _) = tokio::join!(compile_handle, rx_handle); + } + + let (run_handle, run_code_rx, end_tx) = container.run_code().await; + + let buf = Arc::new(Mutex::new(String::default())); + let b = Arc::clone(&buf); + let rx_handle = tokio::spawn(async move { + while let Ok(Some(msg)) = run_code_rx.recv() { + print!("{}", msg); + *b.lock().unwrap() += &msg.to_string(); + } + }); + + tokio::spawn(async move { + sleep_until(Instant::now() + Duration::from_secs(10)).await; + end_tx.send(()).unwrap(); + }); + + let (_, _) = tokio::join!(run_handle, rx_handle); + + message + .edit( + &context.http, + EditMessage::new().content(format!( + "Result\n```{}\n```", + buf.lock().unwrap().replace("@", "\\@") + )), + ) + .await + .unwrap(); + + container.stop().await; + } + } + } + + async fn interaction_create(&self, context: Context, interaction: Interaction) { + if let Interaction::ApplicationCommand(command_interaction) = interaction.clone() { + if command_interaction.data.name == "docker" { + if command_interaction.data.options()[0].name == "ps" { + let list = docker_ps().await; + + let list: Vec = list + .into_iter() + .filter(|p| p.state.clone().unwrap() == "running") + .map(|f| { + String::from(format!("{} {}", f.names.unwrap()[0], f.image.unwrap())) + }) + .collect(); + + let list = list.join("\n"); + + command_interaction + .create_interaction_response( + &context.http, + CreateInteractionResponse::Message( + CreateInteractionResponseMessage::new().content(list), + ), + ) + .await + .unwrap(); + } + } + } + } + + async fn ready(&self, context: Context, _: Ready) { + info!("Ready."); + let docker_command = CreateApplicationCommand::new("docker") + .kind(CommandType::ChatInput) + .description("docker command") + .add_option(CreateApplicationCommandOption::new( + CommandOptionType::SubCommand, + "ps", + "docker ps", + )) + .add_option(CreateApplicationCommandOption::new( + CommandOptionType::SubCommand, + "help", + "docker help command", + )); + + Command::create_global_application_command(&context.http, docker_command).await.unwrap(); + } +} diff --git a/src/language.rs b/src/language.rs new file mode 100644 index 0000000..3d5e46e --- /dev/null +++ b/src/language.rs @@ -0,0 +1,46 @@ +use bollard::{container::Config, service::HostConfig}; +use serde::{Serialize, Deserialize}; + +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct Language { + pub name: String, + pub code: Vec, + pub extension: String, + pub path: String, + pub run_command: String, + pub compile_command: Option, + pub image: String +} + +impl Language { + pub fn get_path(&self, file_name: String) -> String { + self.path.clone().replace("{file}", &file_name) + } + + pub fn get_run_command(&self, file_name: String) -> String { + self.run_command.clone().replace("{file}", &file_name) + } + + pub fn get_compile_command(&self, file_name: String) -> Option { + if let Some(compile) = self.compile_command.clone() { + Some(compile.replace("{file}", &file_name)) + } else { + None + } + } + + pub fn get_container_option(&self) -> Config<&str> { + Config { + image: Some(&self.image), + tty: Some(true), + cmd: Some(vec!["/bin/sh"]), + network_disabled: Some(true), + stop_timeout: Some(30), + host_config: Some(HostConfig { + memory: Some(1024 * 1024 * 1024), + ..Default::default() + }), + ..Default::default() + } + } +} \ No newline at end of file diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..16f2646 --- /dev/null +++ b/src/main.rs @@ -0,0 +1,80 @@ +mod config; +mod docker; +mod event_handler; +mod language; + +use std::{env, fs::File, io::Read, sync::{Arc, Mutex}}; + +use config::Config; +use event_handler::Handler; +use serenity::{ + framework::StandardFramework, http::Http, + prelude::{GatewayIntents, TypeMapKey}, Client, +}; + +struct ConfigStorage; + +impl TypeMapKey for ConfigStorage { + type Value = Arc>; +} + +fn load_config() -> Option { + if let Ok(mut config_file) = File::open("config.yaml") { + let mut buf = String::default(); + config_file.read_to_string(&mut buf).unwrap(); + let config = serde_yaml::from_str::(&buf); + + if let Ok(config) = config { + Some(config) + } else { + None + } + } else { + None + } +} + +#[tokio::main] +async fn main() -> Result<(), ()> { + env::set_var("RUST_LOG", "info"); + env_logger::init(); + + let config = load_config().unwrap(); + + let token = &config.token; + let http = Http::new(token); + + let bot_id = match http.get_current_user().await { + Ok(bot_id) => bot_id.id, + Err(why) => panic!("Could not access the bot id: {:?}", why), + }; + + let framework = StandardFramework::new(); + framework.configure(|c| { + c.with_whitespace(true) + .on_mention(Some(bot_id)) + .prefix(config.prefix.clone()) + }); + + let mut client = Client::builder( + &token, + GatewayIntents::MESSAGE_CONTENT + | GatewayIntents::GUILD_MESSAGES + | GatewayIntents::DIRECT_MESSAGES, + ) + .framework(framework) + .event_handler(Handler) + .await + .expect("Err creating client"); + + { + let mut data = client.data.write().await; + data.insert::(Arc::new(Mutex::new(config))); + } + + if let Err(why) = client.start().await { + println!("Client error: {:?}", why); + } + + Ok(()) +}