├── .github ├── ISSUE_TEMPLATE │ ├── bug_report.md │ └── feature_request.md └── workflows │ └── ci.yml ├── .gitignore ├── .gitmodules ├── .vscode └── settings.json ├── CODE_OF_CONDUCT.md ├── Cargo.toml ├── LICENSE ├── README.md ├── assets └── assyst.svg ├── assyst-cache ├── Cargo.toml ├── README.md └── src │ ├── caches │ ├── guilds.rs │ └── mod.rs │ └── main.rs ├── assyst-common ├── Cargo.toml └── src │ ├── cache │ └── mod.rs │ ├── config │ ├── config.rs │ └── mod.rs │ ├── eval.rs │ ├── lib.rs │ ├── macros.rs │ ├── metrics_handler.rs │ ├── pipe │ ├── mod.rs │ └── pipe_server.rs │ └── util │ ├── discord.rs │ ├── filetype.rs │ ├── mod.rs │ ├── process.rs │ ├── rate_tracker.rs │ ├── regex.rs │ └── table.rs ├── assyst-core ├── Cargo.toml └── src │ ├── assyst.rs │ ├── bad_translator.rs │ ├── command │ ├── arguments.rs │ ├── autocomplete.rs │ ├── componentctxt.rs │ ├── errors.rs │ ├── flags.rs │ ├── fun │ │ ├── colour.rs │ │ ├── mod.rs │ │ └── translation.rs │ ├── group.rs │ ├── image │ │ ├── audio.rs │ │ ├── bloom.rs │ │ ├── caption.rs │ │ ├── makesweet.rs │ │ ├── mod.rs │ │ ├── randomize.rs │ │ └── speechbubble.rs │ ├── messagebuilder.rs │ ├── misc │ │ ├── btchannel.rs │ │ ├── help.rs │ │ ├── mod.rs │ │ ├── prefix.rs │ │ ├── remind.rs │ │ ├── run.rs │ │ ├── stats.rs │ │ └── tag.rs │ ├── mod.rs │ ├── registry.rs │ ├── services │ │ ├── cooltext.rs │ │ ├── download.rs │ │ └── mod.rs │ └── source.rs │ ├── command_ratelimits.rs │ ├── downloader.rs │ ├── gateway_handler │ ├── event_handlers │ │ ├── channel_update.rs │ │ ├── entitlement_create.rs │ │ ├── entitlement_delete.rs │ │ ├── entitlement_update.rs │ │ ├── guild_create.rs │ │ ├── guild_delete.rs │ │ ├── guild_update.rs │ │ ├── interaction_create.rs │ │ ├── message_create.rs │ │ ├── message_delete.rs │ │ ├── message_update.rs │ │ ├── mod.rs │ │ └── ready.rs │ ├── incoming_event.rs │ ├── message_parser │ │ ├── error.rs │ │ ├── mod.rs │ │ ├── parser.rs │ │ └── preprocess.rs │ ├── mod.rs │ └── reply.rs │ ├── main.rs │ ├── persistent_cache_handler │ └── mod.rs │ ├── replies.rs │ ├── rest │ ├── audio_identification.rs │ ├── bad_translation.rs │ ├── charinfo.rs │ ├── cooltext.rs │ ├── eval.rs │ ├── filer.rs │ ├── identify.rs │ ├── mod.rs │ ├── patreon.rs │ ├── r34.rs │ ├── rest_cache_handler.rs │ ├── rust.rs │ ├── top_gg.rs │ └── web_media_download.rs │ └── task │ ├── mod.rs │ └── tasks │ ├── get_premium_users.rs │ ├── mod.rs │ ├── refresh_entitlements.rs │ ├── reminders.rs │ └── top_gg_stats.rs ├── assyst-database ├── Cargo.toml ├── README.md └── src │ ├── cache.rs │ ├── lib.rs │ └── model │ ├── active_guild_premium_entitlement.rs │ ├── badtranslator_channel.rs │ ├── badtranslator_messages.rs │ ├── colour_role.rs │ ├── command_usage.rs │ ├── free_tier_2_requests.rs │ ├── global_blacklist.rs │ ├── guild_disabled_command.rs │ ├── mod.rs │ ├── prefix.rs │ ├── reminder.rs │ ├── tag.rs │ └── user_votes.rs ├── assyst-flux-iface ├── Cargo.toml ├── README.md └── src │ ├── flux_request.rs │ ├── jobs.rs │ ├── lib.rs │ └── limits.rs ├── assyst-gateway ├── Cargo.toml ├── README.md ├── assets │ └── gateway.png └── src │ └── main.rs ├── assyst-proc-macro ├── Cargo.toml ├── README.md └── src │ └── lib.rs ├── assyst-string-fmt ├── Cargo.toml ├── README.md └── src │ ├── ansi.rs │ ├── lib.rs │ └── markdown.rs ├── assyst-tag ├── Cargo.toml ├── fuzz │ ├── .gitignore │ ├── Cargo.toml │ └── fuzz_targets │ │ └── fuzz_target_1.rs └── src │ ├── context.rs │ ├── errors.rs │ ├── lib.rs │ ├── parser.rs │ └── subtags.rs ├── assyst-webserver ├── Cargo.toml └── src │ └── lib.rs ├── config.template.toml ├── run.sh ├── rust-toolchain └── rustfmt.toml /.github/ISSUE_TEMPLATE/bug_report.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Bug report 3 | about: Create a report to help us improve 4 | title: '' 5 | labels: bug 6 | assignees: '' 7 | 8 | --- 9 | 10 | **Describe the bug** 11 | A clear and concise description of what the bug is. 12 | 13 | **To Reproduce** 14 | Steps to reproduce the behavior: 15 | 16 | **Expected behavior** 17 | A clear and concise description of what you expected to happen. 18 | 19 | **Screenshots** 20 | If applicable, add screenshots to help explain your problem. 21 | 22 | **Additional context** 23 | Add any other context about the problem here. 24 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/feature_request.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Feature request 3 | about: Suggest an idea for the bot 4 | title: '' 5 | labels: enhancement 6 | assignees: '' 7 | 8 | --- 9 | 10 | **What would you like to see added to Assyst?** 11 | Please describe in detail what you'd like to see and why. 12 | 13 | **Additional context** 14 | Add any other context or screenshots about the feature request here. 15 | -------------------------------------------------------------------------------- /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | name: Nightly Build/test 2 | 3 | on: 4 | push: 5 | pull_request: 6 | 7 | env: 8 | CARGO_TERM_COLOR: always 9 | 10 | jobs: 11 | build_and_test: 12 | name: Rust project - latest 13 | runs-on: ubuntu-latest 14 | strategy: 15 | matrix: 16 | toolchain: 17 | - nightly 18 | steps: 19 | - uses: actions/checkout@v3 20 | - run: rustup update ${{ matrix.toolchain }} && rustup default ${{ matrix.toolchain }} 21 | - run: cargo build --verbose 22 | - run: cargo test --verbose 23 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | debug/ 2 | target/ 3 | Cargo.lock 4 | config.toml 5 | **/*.rs.bk 6 | *.pdb 7 | .patreon_refresh 8 | **/*.swp -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "flux"] 2 | path = flux 3 | url = git@github.com:Jacherr/flux.git 4 | -------------------------------------------------------------------------------- /.vscode/settings.json: -------------------------------------------------------------------------------- 1 | { 2 | "rust-analyzer.showUnlinkedFileNotification": false, 3 | "editor.defaultFormatter": "rust-lang.rust-analyzer", 4 | "editor.formatOnSave": true, 5 | "[toml]": { 6 | "editor.defaultFormatter": "tamasfe.even-better-toml" 7 | }, 8 | "rust-analyzer.check.command": "clippy", 9 | "terminal.integrated.env.linux": { 10 | "GTK_PATH": null, 11 | "GIO_MODULE_DIR": null, 12 | }, 13 | "[json]": { 14 | "editor.defaultFormatter": "vscode.json-language-features" 15 | }, 16 | } -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | cargo-features = ["codegen-backend"] 2 | 3 | [workspace] 4 | members = [ 5 | "assyst-cache", 6 | "assyst-common", 7 | "assyst-core", 8 | "assyst-database", 9 | "assyst-flux-iface", 10 | "assyst-gateway", 11 | "assyst-proc-macro", 12 | "assyst-string-fmt", 13 | "assyst-tag", 14 | "assyst-webserver", 15 | ] 16 | exclude = ["flux"] 17 | resolver = "2" 18 | 19 | [workspace.lints.clippy] 20 | uninlined_format_args = "warn" 21 | redundant_clone = "warn" 22 | too_long_first_doc_paragraph = "allow" 23 | 24 | [workspace.dependencies] 25 | anyhow = "1.0.75" 26 | serde = { version = "1.0.123", features = ["derive"] } 27 | tokio = { version = "1.34.0", features = ["full"] } 28 | tracing = "0.1.37" 29 | twilight-gateway = { git = "https://github.com/twilight-rs/twilight" } 30 | twilight-http = { git = "https://github.com/twilight-rs/twilight", features = [ 31 | "rustls-ring", 32 | ] } 33 | twilight-model = { git = "https://github.com/twilight-rs/twilight" } 34 | twilight-util = { git = "https://github.com/twilight-rs/twilight", features = [ 35 | "builder", 36 | ] } 37 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 James Croucher 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 |
2 | 3 |
4 | 5 | #
Assyst
6 | 7 |
8 | 9 | ![Discord](https://img.shields.io/discord/1099115731301449758?color=7289DA) 10 | ![GitHub](https://img.shields.io/github/license/jacherr/assyst2) 11 | [![Discord Bots](https://top.gg/api/widget/servers/571661221854707713.svg?noavatar=true)](https://top.gg/bot/571661221854707713) 12 | [![Discord Bots](https://top.gg/api/widget/status/571661221854707713.svg?noavatar=true)](https://top.gg/bot/571661221854707713) 13 | 14 |
15 | 16 | Assyst is a multi-purpose Discord bot with a focus on image processing and manipulation, custom commands via a tag parser, and other unique features. A more detailed overview of the Assyst feature-set can be located on the [Top.gg listing page for Assyst](https://top.gg/bot/571661221854707713). 17 | 18 | Assyst is powered by a custom image editing service called `Flux`. Flux is available [here](https://github.com/jacherr/flux). It provides image and video editing features via the command-line, and is a submodule of this repository in order to facilitate single-step deploys. Refer to the Flux README for more information. 19 | 20 | Assyst is split into a number of separate crates, as described below. 21 | 22 | ## Binaries 23 | - assyst-core: Main command-handling process. Also contains logic for the parsing of message-based commands. 24 | - assyst-gateway: Connects to the Discord WebSocket gateway to receive messages, which are then forwarded to assyst-core for processing. 25 | - assyst-cache: Independent cache process designed to hold some caching information. 26 | 27 | ## Libraries 28 | - assyst-common: Utilities, structures, and functions shared throughout the entire Assyst ecosystem. 29 | - assyst-tag: Tag parser and handler. 30 | - assyst-database: Interfaces with PostgreSQL, for database purposes. 31 | - assyst-webserver: Web server designed to handle webhooking, such as vote processing for Discord bot list websites, as well as Prometheus metrics. 32 | - assyst-proc-macro: General purpose [procedural macros] (currently just a macro for command setup) 33 | - assyst-flux-iface: Basic wrapper over Flux for ease of use. 34 | - assyst-string-fmt: String parsing and formatting utilities. 35 | 36 | ### Note: assyst-core will likely be split into more, smaller, crates in the future. 37 | 38 | [Procedural macros]: https://doc.rust-lang.org/reference/procedural-macros.html 39 | 40 | For more information on each crate, refer to the README.md file for the crate. 41 | 42 | Each binary is ran as an independent process on the same host machine. Each binary communicates through the use of Unix-like pipes. For more information, please refer to the README.md file for the relevant crate. 43 | 44 | ## Contributing 45 | 46 | All contributions - both issues and pull requests - are greatly appreciated. Contributions are done on a fairly loose basis. The easiest way to begin contributing is to first understand the structure of Assyst - this can be done initially by understanding all individual crates by reading their READMEs. If you have any questions, feel free to open an issue. All issues are free to be tackled by anyone. 47 | 48 | ## Self-hosting 49 | 50 | Self-hosting is not yet supported for this version of Assyst, since it is not yet considered production-ready. Self-hosting may be supported with release 1.0.0. \ 51 | However, for completeness, the entire tech stack of Assyst is as follows: 52 | - Rust, as well as Cargo for building. 53 | - PostgreSQL. Database format TBA. 54 | - Flux, which has [its own set of requirements](https://github.com/jacherr/flux?tab=readme-ov-file#prerequisites) 55 | - [youtube-dlp](https://github.com/yt-dlp/yt-dlp) 56 | - fake-eval service (currently closed source). 57 | - CDN (filer) service (currently closed course). 58 | - nginx, for top.gg and prometheus webserver. 59 | - Optionally, Grafana and Prometheus for graphs and logging. A template for this may be made available eventually. If you would like it, open an issue. 60 | 61 | ## Acknowledgements 62 | 63 | Special thanks to [y21](https://github.com/y21) and [Mina](https://github.com/trueharuu) for their invaluable help and contributions towards this version of Assyst. \ 64 | Thank you to the team developing [cobalt.tools](https://github.com/imputnet/cobalt) for creating such a versatile and easy-to-use downloading tool. \ 65 | Thank you to the countless developers of the libraries and programs powering both Assyst and Flux, in particular: 66 | - [Tokio](https://github.com/tokio/tokio) - the asynchronous runtime that Assyst uses, 67 | - [Twilight](https://github.com/twilight-rs/twilight) - the Discord API library that Assyst communicates to Discord with, 68 | - [Image](https://github.com/image-rs/image) - the primary library providing image decoding, encoding, and editing functionality to Flux, 69 | - [FFmpeg](https://ffmpeg.org) - simply the best multimedia processing tool ever made, 70 | - [gegl](https://gegl.org) - providing a bunch of handy image manipulation tools. 71 | 72 | -------------------------------------------------------------------------------- /assets/assyst.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | -------------------------------------------------------------------------------- /assyst-cache/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "assyst-cache" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | anyhow = { workspace = true } 10 | tokio = { workspace = true } 11 | tracing = { workspace = true } 12 | assyst-common = { path = "../assyst-common" } 13 | dashmap = "5.5.3" 14 | rustc-hash = "1.1.0" 15 | 16 | [lints] 17 | workspace = true 18 | -------------------------------------------------------------------------------- /assyst-cache/README.md: -------------------------------------------------------------------------------- 1 | # assyst-cache 2 | 3 | assyst-cache is a relatively simple crate which holds maps and sets of data which may be undesirable to lose in the event of a restart of assyst-core, but are not important enoughto necessitate a table in the database for them. 4 | 5 | One good example of this is the list of guild IDs Assyst is in - this information is *only* sent when the Discord WebSocket gateway connects to Assyst, and so if assyst-core held this information, it would be lost and impossible to retrieve without an (expensive) restart of the gateway. As a result, this crate holds this data so after assyst-core is restarted, it is held and can still be accessed. 6 | 7 | For more information on each cache, refer to the doc comment for the cache. -------------------------------------------------------------------------------- /assyst-cache/src/caches/guilds.rs: -------------------------------------------------------------------------------- 1 | use std::hash::BuildHasherDefault; 2 | 3 | use assyst_common::cache::{GuildCreateData, GuildDeleteData, ReadyData}; 4 | use dashmap::DashSet; 5 | 6 | /// The cache of all guilds Assyst is part of. When a shard readies up, it receives a list of all 7 | /// guild IDs that shard is responsible for, which are all cached here. In addition, when a shard is 8 | /// ready, it will receive `GUILD_CREATE` events for every guild that shard is responsible for. This 9 | /// cache allows to differentiate between `GUILD_CREATE` events sent as part of this procedure (since 10 | /// they were also part of the READY event) and legitimate `GUILD_CREATEs` fired as a result of Assyst 11 | /// joining a new guild post-ready. 12 | pub struct GuildCache { 13 | ids: DashSet>, 14 | } 15 | impl GuildCache { 16 | pub fn new() -> GuildCache { 17 | GuildCache { 18 | ids: DashSet::with_hasher(BuildHasherDefault::::default()), 19 | } 20 | } 21 | 22 | /// Handles a READY event, caching its guilds. Returns the number of newly cached guilds. 23 | pub fn handle_ready_event(&mut self, event: ReadyData) -> u64 { 24 | let mut new_guilds = 0; 25 | 26 | for guild in event.guilds { 27 | if self.ids.insert(guild) { 28 | new_guilds += 1; 29 | }; 30 | } 31 | 32 | new_guilds 33 | } 34 | 35 | /// Handles a `GUILD_CREATE`. This method returns a bool which states if this guild is new or not. 36 | /// A new guild is one that was not received during the start-up of the gateway connection. 37 | pub fn handle_guild_create_event(&mut self, event: GuildCreateData) -> bool { 38 | self.ids.insert(event.id) 39 | } 40 | 41 | /// Handles a `GUILD_DELETE`. This method returns a bool which states if the bot was actually 42 | /// kicked from this guild. 43 | pub fn handle_guild_delete_event(&mut self, event: GuildDeleteData) -> bool { 44 | !event.unavailable && self.ids.remove(&event.id).is_some() 45 | } 46 | 47 | pub fn size(&self) -> u64 { 48 | self.ids.len() as u64 49 | } 50 | } 51 | -------------------------------------------------------------------------------- /assyst-cache/src/caches/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod guilds; 2 | -------------------------------------------------------------------------------- /assyst-cache/src/main.rs: -------------------------------------------------------------------------------- 1 | use assyst_common::cache::{CacheJob, CacheResponse}; 2 | use assyst_common::ok_or_break; 3 | use assyst_common::pipe::pipe_server::PipeServer; 4 | use assyst_common::pipe::CACHE_PIPE_PATH; 5 | use assyst_common::util::tracing_init; 6 | use tracing::{debug, info, warn}; 7 | 8 | use crate::caches::guilds::GuildCache; 9 | 10 | mod caches; 11 | 12 | #[tokio::main] 13 | async fn main() { 14 | tracing_init(); 15 | 16 | let mut guild_cache = GuildCache::new(); 17 | 18 | let mut pipe_server = PipeServer::listen(CACHE_PIPE_PATH).unwrap(); 19 | info!("Awaiting connection from assyst-core"); 20 | loop { 21 | let mut stream = pipe_server.accept_connection().await.unwrap(); 22 | info!("Connection received from assyst-core"); 23 | loop { 24 | let job = ok_or_break!(stream.read_object::().await); 25 | 26 | debug!("Handling job: {}", job); 27 | 28 | let result = match job { 29 | CacheJob::HandleReady(event) => { 30 | CacheResponse::NewGuildsFromReady(guild_cache.handle_ready_event(event)) 31 | }, 32 | CacheJob::HandleGuildCreate(event) => { 33 | CacheResponse::ShouldHandleGuildCreate(guild_cache.handle_guild_create_event(event)) 34 | }, 35 | CacheJob::HandleGuildDelete(event) => { 36 | CacheResponse::ShouldHandleGuildDelete(guild_cache.handle_guild_delete_event(event)) 37 | }, 38 | CacheJob::GetGuildCount => CacheResponse::TotalGuilds(guild_cache.size()), 39 | }; 40 | 41 | ok_or_break!(stream.write_object(result).await); 42 | } 43 | warn!("Connection to assyst-core lost, awaiting reconnection"); 44 | } 45 | } 46 | -------------------------------------------------------------------------------- /assyst-common/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "assyst-common" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | assyst-string-fmt = { path = "../assyst-string-fmt" } 10 | anyhow = { workspace = true } 11 | bincode = "1.3.3" 12 | serde = { workspace = true } 13 | tokio = { workspace = true } 14 | toml = "0.8.8" 15 | lazy_static = "1.4.0" 16 | tracing = { workspace = true } 17 | assyst-database = { path = "../assyst-database" } 18 | twilight-http = { workspace = true } 19 | twilight-model = { workspace = true } 20 | prometheus = "0.13.3" 21 | regex = "1.4.3" 22 | reqwest = { version = "0.11.24" } 23 | tracing-subscriber = { version = "0.3.16", features = ["time", "env-filter"] } 24 | time = { version = "0.3.31", features = ["macros"] } 25 | num_cpus = "1.16.0" 26 | rayon = "1.8.1" 27 | rand = "0.8.5" 28 | 29 | [lints] 30 | workspace = true 31 | -------------------------------------------------------------------------------- /assyst-common/src/cache/mod.rs: -------------------------------------------------------------------------------- 1 | use std::fmt::Display; 2 | 3 | use serde::{Deserialize, Serialize}; 4 | use tokio::sync::oneshot::Sender; 5 | use twilight_model::gateway::payload::incoming::{GuildCreate, GuildDelete, Ready}; 6 | 7 | /// A cache job, along with a transmitter to send the response back to the calling thread. 8 | pub type CacheJobSend = (Sender, CacheJob); 9 | 10 | /// The different jobs that the cache needs to handle. 11 | #[derive(Serialize, Deserialize, Debug)] 12 | pub enum CacheJob { 13 | /// Storing data from a `GUILD_CREATE` event. 14 | HandleGuildCreate(GuildCreateData), 15 | /// Storing data from a `GUILD_DELETE` event. 16 | HandleGuildDelete(GuildDeleteData), 17 | /// Storing data from a READY event. 18 | HandleReady(ReadyData), 19 | GetGuildCount, 20 | } 21 | impl Display for CacheJob { 22 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 23 | match self { 24 | Self::HandleReady(x) => write!(f, "HandleReady ({} guilds)", x.guilds.len()), 25 | Self::HandleGuildCreate(x) => write!(f, "HandleGuildCreate (ID: {})", x.id), 26 | Self::HandleGuildDelete(x) => write!(f, "HandleGuildDelete (ID: {})", x.id), 27 | Self::GetGuildCount => f.write_str("GetGuildCount"), 28 | } 29 | } 30 | } 31 | 32 | #[derive(Serialize, Deserialize, Debug)] 33 | pub struct ReadyData { 34 | pub guilds: Vec, 35 | } 36 | impl From for ReadyData { 37 | fn from(value: Ready) -> Self { 38 | ReadyData { 39 | guilds: value.guilds.iter().map(|x| x.id.get()).collect::>(), 40 | } 41 | } 42 | } 43 | 44 | #[derive(Serialize, Deserialize, Debug)] 45 | pub struct GuildCreateData { 46 | pub id: u64, 47 | pub name: String, 48 | pub members: Option, 49 | } 50 | impl From for GuildCreateData { 51 | fn from(value: GuildCreate) -> Self { 52 | match value { 53 | GuildCreate::Available(g) => GuildCreateData { 54 | id: g.id.get(), 55 | name: g.name.clone(), 56 | members: g.member_count, 57 | }, 58 | GuildCreate::Unavailable(g) => GuildCreateData { 59 | id: g.id.get(), 60 | name: String::new(), 61 | members: None, 62 | }, 63 | } 64 | } 65 | } 66 | 67 | #[derive(Serialize, Deserialize, Debug)] 68 | pub struct GuildDeleteData { 69 | pub id: u64, 70 | pub unavailable: bool, 71 | } 72 | impl From for GuildDeleteData { 73 | fn from(value: GuildDelete) -> Self { 74 | GuildDeleteData { 75 | id: value.id.get(), 76 | unavailable: value.unavailable.unwrap_or(false), 77 | } 78 | } 79 | } 80 | 81 | pub type CacheResponseSend = anyhow::Result; 82 | 83 | /// All the responses the cache can send back. Usually it is a 1-1 relation between a `CacheJob` 84 | /// variant and `CacheResponse` variant. 85 | #[derive(Serialize, Deserialize, Debug)] 86 | pub enum CacheResponse { 87 | /// Whether Assyst should handle a `GUILD_CREATE` event. False if this guild is coming back from 88 | /// unavailable, or if this guild has already been cached. 89 | ShouldHandleGuildCreate(bool), 90 | /// Whether Assyst should handle a `GUILD_DELETE` event. False if this guild went unavailable, or 91 | /// if it was not in the cache. 92 | ShouldHandleGuildDelete(bool), 93 | /// The amount of new guilds Assyst receives when a shard enters a READY state. Some guilds may 94 | /// be duplicated, which is why this number may differ from the length of the guilds array in 95 | /// this event. 96 | NewGuildsFromReady(u64), 97 | /// Total number of cached guilds. 98 | TotalGuilds(u64), 99 | } 100 | -------------------------------------------------------------------------------- /assyst-common/src/config/config.rs: -------------------------------------------------------------------------------- 1 | // See config.toml for information on the variables here. 2 | 3 | use serde::Deserialize; 4 | 5 | #[derive(Deserialize)] 6 | pub struct AssystConfig { 7 | pub bot_id: u64, 8 | pub urls: Urls, 9 | pub authentication: Authentication, 10 | pub database: Database, 11 | pub prefix: Prefixes, 12 | pub logging_webhooks: LoggingWebhooks, 13 | pub dev: DevAttributes, 14 | pub entitlements: Entitlements, 15 | } 16 | 17 | #[derive(Deserialize)] 18 | pub struct Entitlements { 19 | pub premium_server_sku_id: u64, 20 | } 21 | 22 | #[derive(Deserialize, Clone)] 23 | pub struct CobaltApiInstance { 24 | pub url: String, 25 | pub key: String, 26 | pub primary: Option, 27 | } 28 | 29 | #[derive(Deserialize, Clone)] 30 | pub struct Urls { 31 | pub proxy: Vec, 32 | pub filer: String, 33 | pub eval: String, 34 | pub bad_translation: String, 35 | pub cobalt_api: Vec, 36 | } 37 | 38 | #[derive(Deserialize)] 39 | pub struct Authentication { 40 | pub discord_token: String, 41 | pub patreon_client_secret: String, 42 | pub patreon_client_id: String, 43 | pub top_gg_token: String, 44 | pub top_gg_webhook_token: String, 45 | pub top_gg_webhook_port: u16, 46 | pub filer_key: String, 47 | pub notsoapi: String, 48 | pub rapidapi_token: String, 49 | } 50 | 51 | #[derive(Deserialize)] 52 | pub struct Database { 53 | pub host: String, 54 | pub username: String, 55 | pub password: String, 56 | pub database: String, 57 | pub port: u16, 58 | } 59 | impl Database { 60 | #[must_use] 61 | pub fn to_url(&self) -> String { 62 | format!( 63 | "postgres://{}:{}@{}:{}/{}", 64 | self.username, self.password, self.host, self.port, self.database 65 | ) 66 | } 67 | 68 | #[must_use] 69 | pub fn to_url_safe(&self) -> String { 70 | let mut host = self.host.split('.').take(2).collect::>(); 71 | host.push("###"); 72 | host.push("###"); 73 | 74 | let mut port = self.port.to_string(); 75 | port.replace_range(..3, "..."); 76 | 77 | format!( 78 | "postgres://{}@{}:{}/{}", 79 | self.username, 80 | &host.join("."), 81 | port, 82 | self.database 83 | ) 84 | } 85 | } 86 | 87 | #[derive(Deserialize)] 88 | pub struct Prefixes { 89 | pub default: String, 90 | } 91 | 92 | #[derive(Deserialize)] 93 | pub struct LoggingWebhooks { 94 | pub panic: LoggingWebhook, 95 | pub error: LoggingWebhook, 96 | pub vote: LoggingWebhook, 97 | pub enable_webhooks: bool, 98 | } 99 | 100 | #[derive(Deserialize, Clone)] 101 | pub struct LoggingWebhook { 102 | pub token: String, 103 | pub id: u64, 104 | } 105 | 106 | #[derive(Deserialize)] 107 | pub struct DevAttributes { 108 | pub admin_users: Vec, 109 | pub prefix_override: Option, 110 | pub disable_bad_translator_channels: bool, 111 | pub disable_reminder_check: bool, 112 | pub disable_bot_list_posting: bool, 113 | pub disable_patreon_synchronisation: bool, 114 | pub disable_entitlement_fetching: bool, 115 | pub dev_guild: u64, 116 | pub dev_channel: u64, 117 | pub dev_message: bool, 118 | pub flux_workspace_root_path_override: String, 119 | } 120 | -------------------------------------------------------------------------------- /assyst-common/src/config/mod.rs: -------------------------------------------------------------------------------- 1 | #[allow(clippy::module_inception)] 2 | pub mod config; 3 | 4 | pub static CONFIG_LOCATION: &str = "./config.toml"; 5 | pub static PATREON_REFRESH_LOCATION: &str = "./.patreon_refresh"; 6 | 7 | use lazy_static::lazy_static; 8 | use toml::from_str; 9 | use tracing::info; 10 | 11 | use crate::config::config::AssystConfig; 12 | 13 | lazy_static! { 14 | pub static ref CONFIG: AssystConfig = { 15 | let toml = std::fs::read_to_string(CONFIG_LOCATION).unwrap(); 16 | let config = from_str::(&toml).unwrap(); 17 | info!("Loaded config file {}", CONFIG_LOCATION); 18 | config 19 | }; 20 | } 21 | -------------------------------------------------------------------------------- /assyst-common/src/eval.rs: -------------------------------------------------------------------------------- 1 | use serde::{Deserialize, Serialize}; 2 | 3 | use crate::util::filetype::Type; 4 | 5 | #[derive(Serialize, Deserialize)] 6 | pub struct FakeEvalMessageData { 7 | pub message: M, 8 | pub args: Vec, 9 | } 10 | 11 | #[derive(Serialize)] 12 | pub struct FakeEvalBody { 13 | pub code: String, 14 | pub data: Option>, 15 | } 16 | 17 | #[derive(Deserialize)] 18 | pub struct FakeEvalResponse { 19 | pub message: String, 20 | } 21 | 22 | pub enum FakeEvalImageResponse { 23 | Text(FakeEvalResponse), 24 | Image(Vec, Type), 25 | } 26 | -------------------------------------------------------------------------------- /assyst-common/src/lib.rs: -------------------------------------------------------------------------------- 1 | #![feature(duration_constants, trait_alias)] 2 | 3 | pub mod cache; 4 | pub mod config; 5 | pub mod eval; 6 | pub mod macros; 7 | pub mod metrics_handler; 8 | pub mod pipe; 9 | pub mod util; 10 | -------------------------------------------------------------------------------- /assyst-common/src/macros.rs: -------------------------------------------------------------------------------- 1 | use twilight_http::Client as HttpClient; 2 | use twilight_model::id::marker::WebhookMarker; 3 | use twilight_model::id::Id; 4 | 5 | use crate::config::config::LoggingWebhook; 6 | use crate::config::CONFIG; 7 | 8 | #[macro_export] 9 | macro_rules! ok_or_break { 10 | ($expression:expr) => { 11 | match $expression { 12 | Ok(v) => v, 13 | Err(_) => break, 14 | } 15 | }; 16 | } 17 | 18 | #[macro_export] 19 | macro_rules! ok_or_continue { 20 | ($expression:expr) => { 21 | match $expression { 22 | Ok(v) => v, 23 | Err(_) => continue, 24 | } 25 | }; 26 | } 27 | 28 | #[macro_export] 29 | macro_rules! unwrap_enum_variant { 30 | ($expression:expr, $variant:path) => { 31 | match $expression { 32 | $variant(v) => v, 33 | _ => unreachable!(), 34 | } 35 | }; 36 | } 37 | 38 | #[macro_export] 39 | macro_rules! err { 40 | ($($t:tt)*) => {{ 41 | use $crate::macros::handle_log; 42 | let msg = format!($($t)*); 43 | tracing::error!("{}", &msg); 44 | 45 | handle_log(format!("Error: ```{}```", msg)); 46 | }} 47 | } 48 | 49 | pub fn handle_log(message: String) { 50 | if CONFIG.logging_webhooks.enable_webhooks { 51 | tokio::spawn(async move { 52 | let LoggingWebhook { id, token } = CONFIG.logging_webhooks.error.clone(); 53 | 54 | let client = HttpClient::new(CONFIG.authentication.discord_token.clone()); 55 | 56 | if id == 0 { 57 | tracing::error!("Failed to trigger error webhook: Error webhook ID is 0"); 58 | } else { 59 | let webhook = client 60 | .execute_webhook(Id::::new(id), &token) 61 | .content(&message); 62 | 63 | let _ = webhook 64 | .await 65 | .inspect_err(|e| tracing::error!("Failed to trigger error webhook: {}", e.to_string())); 66 | } 67 | }); 68 | } 69 | } 70 | -------------------------------------------------------------------------------- /assyst-common/src/metrics_handler.rs: -------------------------------------------------------------------------------- 1 | use std::collections::HashMap; 2 | use std::sync::{Arc, Mutex}; 3 | use std::time::Duration; 4 | 5 | use assyst_database::DatabaseHandler; 6 | use prometheus::{register_int_counter, register_int_gauge_vec, IntCounter, IntGaugeVec}; 7 | use tracing::debug; 8 | 9 | use crate::util::process::get_processes_mem_usage; 10 | use crate::util::rate_tracker::RateTracker; 11 | 12 | /// Handler for general metrics, including rate trackers, Prometheus metrics, etc. 13 | pub struct MetricsHandler { 14 | pub cache_sizes: IntGaugeVec, 15 | pub memory_usage: IntGaugeVec, 16 | pub guilds: IntGaugeVec, 17 | pub guilds_rate_tracker: Mutex, 18 | pub events: IntCounter, 19 | pub events_rate_tracker: Mutex, 20 | pub commands: IntCounter, 21 | pub total_commands_rate_tracker: Mutex, 22 | pub individual_commands_rate_trackers: tokio::sync::Mutex>, 23 | pub database_handler: Arc, 24 | } 25 | impl MetricsHandler { 26 | pub fn new(database_handler: Arc) -> anyhow::Result { 27 | Ok(MetricsHandler { 28 | cache_sizes: register_int_gauge_vec!("cache_sizes", "Cache sizes", &["cache"])?, 29 | memory_usage: register_int_gauge_vec!("memory_usage", "Memory usage in MB", &["process"])?, 30 | guilds: register_int_gauge_vec!("guilds", "Total guilds and user installs", &["context"])?, 31 | guilds_rate_tracker: Mutex::new(RateTracker::new(Duration::from_secs(60 * 60))), 32 | events: register_int_counter!("events", "Total number of events")?, 33 | events_rate_tracker: Mutex::new(RateTracker::new(Duration::from_secs(1))), 34 | commands: register_int_counter!("commands", "Total number of commands executed")?, 35 | total_commands_rate_tracker: Mutex::new(RateTracker::new(Duration::from_secs(60))), 36 | individual_commands_rate_trackers: tokio::sync::Mutex::new(HashMap::new()), 37 | database_handler, 38 | }) 39 | } 40 | 41 | pub fn update_cache_size(&self, cache: &str, size: usize) { 42 | self.cache_sizes.with_label_values(&[cache]).set(size as i64); 43 | } 44 | 45 | /// Updates some metrics that are not updated as data comes in. 46 | pub async fn update(&self, user_installs: u64) { 47 | debug!("Collecting prometheus metrics"); 48 | 49 | let database_cache_reader = &self.database_handler; 50 | let prefixes_cache_size = database_cache_reader.cache.get_prefixes_cache_size(); 51 | self.update_cache_size("prefixes", prefixes_cache_size); 52 | self.update_cache_size( 53 | "disabled_commands", 54 | database_cache_reader.cache.get_guild_disabled_commands_size(), 55 | ); 56 | 57 | let memory_usages = get_processes_mem_usage(); 58 | 59 | self.guilds.with_label_values(&["installs"]).set(user_installs as i64); 60 | 61 | for usage in memory_usages { 62 | self.memory_usage 63 | .with_label_values(&[usage.0]) 64 | .set((usage.1 / 1024 / 1024) as i64); 65 | } 66 | } 67 | 68 | pub fn set_user_installs(&self, installs: u64) { 69 | self.guilds.with_label_values(&["installs"]).set(installs as i64); 70 | } 71 | 72 | pub fn add_guilds(&self, guilds: u64) { 73 | self.guilds.with_label_values(&["guilds"]).add(guilds as i64); 74 | } 75 | 76 | pub fn inc_guilds(&self) { 77 | self.guilds_rate_tracker.lock().unwrap().add_sample(); 78 | self.guilds.with_label_values(&["guilds"]).inc(); 79 | } 80 | 81 | pub fn dec_guilds(&self) { 82 | self.guilds_rate_tracker.lock().unwrap().remove_sample(); 83 | self.guilds.with_label_values(&["guilds"]).dec(); 84 | } 85 | 86 | pub fn add_event(&self) { 87 | self.events.inc(); 88 | self.events_rate_tracker.lock().unwrap().add_sample(); 89 | } 90 | 91 | pub fn get_events_rate(&self) -> usize { 92 | self.events_rate_tracker.lock().unwrap().get_rate() 93 | } 94 | 95 | pub fn add_command(&self) { 96 | self.commands.inc(); 97 | self.total_commands_rate_tracker.lock().unwrap().add_sample(); 98 | } 99 | 100 | pub fn get_commands_rate(&self) -> usize { 101 | self.total_commands_rate_tracker.lock().unwrap().get_rate() 102 | } 103 | 104 | pub async fn add_individual_command_usage(&self, command_name: &'static str) { 105 | let mut lock = self.individual_commands_rate_trackers.lock().await; 106 | let entry = lock.get_mut(&command_name); 107 | if let Some(entry) = entry { 108 | entry.add_sample(); 109 | } else { 110 | let mut tracker = RateTracker::new(Duration::from_secs(60 * 60)); 111 | tracker.add_sample(); 112 | lock.insert(command_name, tracker); 113 | } 114 | } 115 | } 116 | -------------------------------------------------------------------------------- /assyst-common/src/pipe/mod.rs: -------------------------------------------------------------------------------- 1 | use std::time::Duration; 2 | 3 | use anyhow::bail; 4 | use bincode::{deserialize, serialize}; 5 | use serde::de::DeserializeOwned; 6 | use serde::Serialize; 7 | use tokio::io::{AsyncReadExt, AsyncWriteExt}; 8 | use tokio::net::UnixStream; 9 | use tokio::time::sleep; 10 | use tracing::warn; 11 | 12 | use crate::util::string_from_likely_utf8; 13 | 14 | pub mod pipe_server; 15 | 16 | pub static GATEWAY_PIPE_PATH: &str = "/tmp/assyst2-gateway-com"; 17 | pub static CACHE_PIPE_PATH: &str = "/tmp/assyst2-cache-com"; 18 | 19 | static POLL_FREQUENCY: Duration = Duration::from_secs(10); 20 | 21 | /// Pipe is a utility class that wraps a [`UnixStream`], providing helper functions for easy reading 22 | /// and writing of serde-Serializable types via Bincode. 23 | pub struct Pipe { 24 | pub stream: UnixStream, 25 | } 26 | impl Pipe { 27 | /// Connect to a specific file descriptor. 28 | pub async fn connect(pipe_location: &str) -> anyhow::Result { 29 | let stream = UnixStream::connect(pipe_location).await?; 30 | Ok(Pipe { stream }) 31 | } 32 | 33 | /// Repeatedly attempt to connect to a specific file descriptor, until a maximum retry threshold 34 | /// is reached. 35 | pub async fn poll_connect(pipe_location: &str, limit: Option) -> anyhow::Result { 36 | let mut attempts = 0; 37 | 38 | let pipe: Pipe = loop { 39 | let pipe = Pipe::connect(pipe_location).await; 40 | if let Ok(p) = pipe { 41 | break p; 42 | } else if let Err(e) = pipe { 43 | attempts += 1; 44 | warn!( 45 | "{}: connection failed ({}/{:?}): {}", 46 | pipe_location, 47 | attempts, 48 | limit, 49 | e.to_string() 50 | ); 51 | if let Some(l) = limit { 52 | if attempts >= l { 53 | bail!("timed out waiting for connection"); 54 | } 55 | } 56 | sleep(POLL_FREQUENCY).await; 57 | } 58 | }; 59 | 60 | Ok(pipe) 61 | } 62 | 63 | pub fn new(stream: UnixStream) -> Self { 64 | Pipe { stream } 65 | } 66 | 67 | /// Read a Bincode-deserializable object from this stream. 68 | /// 69 | /// This function will return an Err if the stream is prematurely closed, or if Bincode is not 70 | /// able to deserialize the data to the specified type. 71 | pub async fn read_object(&mut self) -> anyhow::Result { 72 | let len = self.stream.read_u32().await?; 73 | let mut data = vec![0u8; len as usize]; 74 | self.stream.read_exact(&mut data).await?; 75 | Ok(deserialize::(&data)?) 76 | } 77 | 78 | /// Read a UTF8-encoded String from this stream. 79 | /// 80 | /// Note: this function heavily favors the "likely UTF-8" case and will be worse 81 | /// for invalid UTF-8 (see [`string_from_likely_utf8`]). 82 | /// This function will return an Err if the stream is prematurely closed. 83 | pub async fn read_string(&mut self) -> anyhow::Result { 84 | let len = self.stream.read_u32().await?; 85 | let mut data = vec![0u8; len as usize]; 86 | self.stream.read_exact(&mut data).await?; 87 | Ok(string_from_likely_utf8(data)) 88 | } 89 | 90 | /// Write a Bincode-serializable object to this stream. 91 | /// 92 | /// This function will return an Err if the stream is prematurely closed, or if Bincode is not 93 | /// able to serialize the data to the specified type. 94 | pub async fn write_object(&mut self, obj: T) -> anyhow::Result<()> { 95 | let buffer = serialize(&obj)?; 96 | debug_assert!(u32::try_from(buffer.len()).is_ok(), "attempted to write more than 4 GB"); 97 | self.stream.write_u32(buffer.len() as u32).await?; 98 | self.stream.write_all(&buffer).await?; 99 | Ok(()) 100 | } 101 | 102 | /// Write a UTF8-encoded String to this stream. 103 | /// 104 | /// This function will return an Err if the stream is prematurely closed. 105 | pub async fn write_string>(&mut self, obj: T) -> anyhow::Result<()> { 106 | let obj = obj.as_ref(); 107 | debug_assert!(u32::try_from(obj.len()).is_ok(), "attempted to write more than 4 GB"); 108 | self.stream.write_u32(obj.len() as u32).await?; 109 | self.stream.write_all(obj.as_bytes()).await?; 110 | Ok(()) 111 | } 112 | } 113 | -------------------------------------------------------------------------------- /assyst-common/src/pipe/pipe_server.rs: -------------------------------------------------------------------------------- 1 | use std::path::Path; 2 | 3 | use tokio::net::UnixListener; 4 | use tracing::info; 5 | 6 | use super::Pipe; 7 | 8 | /// `PipeServer` is a utility class wrapping [`UnixListener`] that provides utility functions 9 | /// for listening on a specific file descriptor ("pipe") and accepting a connection from it. 10 | pub struct PipeServer { 11 | listener: UnixListener, 12 | } 13 | impl PipeServer { 14 | /// Listen on a specific file descriptor. 15 | pub fn listen(pipe_location: &str) -> anyhow::Result { 16 | if Path::new(pipe_location).exists() { 17 | info!("Deleting old pipe file {}", pipe_location); 18 | std::fs::remove_file(pipe_location)?; 19 | } 20 | 21 | let listener = UnixListener::bind(pipe_location)?; 22 | Ok(PipeServer { listener }) 23 | } 24 | 25 | /// Asynchronously wait for a connection to be recieved from the current listener. 26 | pub async fn accept_connection(&mut self) -> anyhow::Result { 27 | let (stream, _) = self.listener.accept().await?; 28 | let pipe = Pipe::new(stream); 29 | Ok(pipe) 30 | } 31 | } 32 | -------------------------------------------------------------------------------- /assyst-common/src/util/discord.rs: -------------------------------------------------------------------------------- 1 | use anyhow::ensure; 2 | use regex::Regex; 3 | use twilight_http::Client; 4 | use twilight_model::id::marker::{ChannelMarker, GuildMarker}; 5 | use twilight_model::id::Id; 6 | use twilight_model::user::User; 7 | 8 | use super::format_time; 9 | use super::regex::USER_MENTION; 10 | 11 | pub const MAX_TIMESTAMP: u64 = 8640000000000000; 12 | 13 | /// Attempts to resolve a guild's owner's user ID 14 | pub async fn get_guild_owner(http: &Client, guild_id: u64) -> anyhow::Result { 15 | Ok(http 16 | .guild(Id::::new(guild_id)) 17 | .await? 18 | .model() 19 | .await? 20 | .owner_id 21 | .get()) 22 | } 23 | 24 | #[must_use] pub fn get_default_avatar_url(user: &User) -> String { 25 | // Unwrapping discrim parsing is ok, it should never be out of range or non-numeric 26 | let suffix = if user.discriminator == 0 { 27 | // Pomelo users 28 | (user.id.get().wrapping_shr(22) % 6) as u16 29 | } else { 30 | // Legacy 31 | user.discriminator % 5 32 | }; 33 | format!("https://cdn.discordapp.com/embed/avatars/{suffix}.png?size=1024") 34 | } 35 | 36 | #[must_use] pub fn get_avatar_url(user: &User) -> String { 37 | let avatar = match &user.avatar { 38 | Some(av) => av, 39 | None => return get_default_avatar_url(user), 40 | }; 41 | 42 | let ext = if avatar.bytes().starts_with("a_".as_bytes()) { 43 | "gif" 44 | } else { 45 | "png" 46 | }; 47 | 48 | format!( 49 | "https://cdn.discordapp.com/avatars/{}/{}.{}?size=1024", 50 | user.id, avatar, ext 51 | ) 52 | } 53 | 54 | #[must_use] pub fn id_from_mention(word: &str) -> Option { 55 | USER_MENTION 56 | .captures(word) 57 | .and_then(|user_id_capture| user_id_capture.get(1)) 58 | .map(|id| id.as_str()) 59 | .and_then(|id| id.parse::().ok()) 60 | } 61 | 62 | #[must_use] pub fn format_tag(user: &User) -> String { 63 | format!("{}#{}", user.name, user.discriminator) 64 | } 65 | 66 | /// Generates a message link 67 | #[must_use] pub fn message_link(guild_id: u64, channel_id: u64, message_id: u64) -> String { 68 | format!("https://discord.com/channels/{guild_id}/{channel_id}/{message_id}") 69 | } 70 | 71 | /// Generates a DM message link 72 | #[must_use] pub fn dm_message_link(channel_id: u64, message_id: u64) -> String { 73 | format!("https://discord.com/channels/@me/{channel_id}/{message_id}") 74 | } 75 | 76 | /// Attempts to return the timestamp as a Discord timestamp, 77 | /// and falls back to [`format_time`] if Discord were to render it as "Invalid Date" 78 | #[must_use] pub fn format_discord_timestamp(input: u64) -> String { 79 | if input <= MAX_TIMESTAMP { 80 | format!("", input / 1000) 81 | } else { 82 | format_time(input) 83 | } 84 | } 85 | 86 | #[must_use] pub fn user_mention_to_id(s: &str) -> Option { 87 | let mention: Regex = Regex::new(r"(?:<@!?)?(\d{16,20})>?").unwrap(); 88 | 89 | mention 90 | .captures(s) 91 | .and_then(|capture| capture.get(1)) 92 | .map(|id| id.as_str()) 93 | .and_then(|id| id.parse::().ok()) 94 | } 95 | 96 | #[must_use] pub fn channel_mention_to_id(s: &str) -> Option { 97 | let mention: Regex = Regex::new(r"(?:<#)?(\d{16,20})>?").unwrap(); 98 | 99 | mention 100 | .captures(s) 101 | .and_then(|capture| capture.get(1)) 102 | .map(|id| id.as_str()) 103 | .and_then(|id| id.parse::().ok()) 104 | } 105 | 106 | pub async fn is_same_guild(client: &Client, channel_id: u64, guild_id: u64) -> Result { 107 | let ch = client 108 | .channel(Id::::new(channel_id)) 109 | .await? 110 | .model() 111 | .await 112 | .unwrap(); 113 | 114 | let real_guild_id = ch.guild_id.map_or(0, twilight_model::id::Id::get); 115 | 116 | Ok(real_guild_id == guild_id) 117 | } 118 | 119 | pub async fn ensure_same_guild(client: &Client, channel_id: u64, guild_id: u64) -> anyhow::Result<()> { 120 | let is = is_same_guild(client, channel_id, guild_id).await?; 121 | 122 | ensure!(is, "The provided channel is not part of this server."); 123 | Ok(()) 124 | } 125 | -------------------------------------------------------------------------------- /assyst-common/src/util/filetype.rs: -------------------------------------------------------------------------------- 1 | use std::cmp::min; 2 | use std::ops::Range; 3 | 4 | #[derive(Debug, PartialEq)] 5 | pub enum Type { 6 | GIF, 7 | JPEG, 8 | PNG, 9 | WEBP, 10 | MP4, 11 | WEBM, 12 | MP3, 13 | ZIP, 14 | } 15 | impl Type { 16 | #[must_use] pub fn as_str(&self) -> &'static str { 17 | match self { 18 | Type::GIF => "gif", 19 | Type::JPEG => "jpeg", 20 | Type::PNG => "png", 21 | Type::WEBP => "webp", 22 | Type::MP4 => "mp4", 23 | Type::WEBM => "webm", 24 | Type::MP3 => "mp3", 25 | Type::ZIP => "zip", 26 | } 27 | } 28 | #[must_use] pub fn as_mime(&self) -> &'static str { 29 | match self { 30 | Type::GIF => "image/gif", 31 | Type::JPEG => "image/jpeg", 32 | Type::PNG => "image/png", 33 | Type::WEBP => "image/webp", 34 | Type::MP4 => "video/mp4", 35 | Type::WEBM => "video/webm", 36 | Type::MP3 => "audio/mpeg", 37 | Type::ZIP => "application/x-zip", 38 | } 39 | } 40 | #[must_use] pub fn is_video(&self) -> bool { 41 | matches!(self, Type::MP4 | Type::WEBM) 42 | } 43 | } 44 | 45 | const WEBP: [u8; 4] = [87, 69, 66, 80]; 46 | const MP4: [u8; 4] = [0x66, 0x74, 0x79, 0x70]; 47 | 48 | fn bounded_range(start: usize, end: usize, len: usize) -> Range { 49 | min(len, start)..min(len, end) 50 | } 51 | 52 | fn sig(that: &[u8], eq: &[u8]) -> bool { 53 | that[0..std::cmp::min(eq.len(), that.len())].eq(eq) 54 | } 55 | 56 | fn check_webp(that: &[u8]) -> bool { 57 | let bytes_offset_removed = &that[bounded_range(8, 12, that.len())]; 58 | sig(bytes_offset_removed, &WEBP) 59 | } 60 | 61 | fn check_mp4(that: &[u8]) -> bool { 62 | let bytes_offset_removed = &that[bounded_range(4, 8, that.len())]; 63 | sig(bytes_offset_removed, &MP4) 64 | } 65 | 66 | #[must_use] pub fn get_sig(buf: &[u8]) -> Option { 67 | match buf { 68 | [71, 73, 70, ..] => Some(Type::GIF), 69 | [255, 216, 255, ..] => Some(Type::JPEG), 70 | [137, 80, 78, 71, 13, 10, 26, 10, ..] => Some(Type::PNG), 71 | [0x1A, 0x45, 0xDF, 0xA3, ..] => Some(Type::WEBM), 72 | [0x49, 0x44, 0x33, ..] /* ID3 tagged */ | [0xff, 0xfb, ..] /* untagged */ => Some(Type::MP3), 73 | [0x50, 0x4b, ..] => Some(Type::ZIP), 74 | _ if check_webp(buf) => Some(Type::WEBP), 75 | _ if check_mp4(buf) => Some(Type::MP4), 76 | _ => None, 77 | } 78 | } 79 | -------------------------------------------------------------------------------- /assyst-common/src/util/rate_tracker.rs: -------------------------------------------------------------------------------- 1 | use std::time::Duration; 2 | 3 | use tokio::time::Instant; 4 | 5 | /// Struct to allow the tracking of how fast a value increases, or how fast a state changes. 6 | /// 7 | /// For example, can be used to determine how frequently a command is ran over a time period, 8 | /// or the rate of events being received. 9 | pub struct RateTracker { 10 | tracking_length: Duration, 11 | samples: Vec, 12 | } 13 | impl RateTracker { 14 | #[must_use] pub fn new(tracking_length: Duration) -> RateTracker { 15 | RateTracker { 16 | tracking_length, 17 | samples: vec![], 18 | } 19 | } 20 | 21 | /// Removes all samples from this tracker which are older than the tracking length. 22 | pub fn remove_expired_samples(&mut self) { 23 | self.samples 24 | .retain(|x| Instant::now().duration_since(*x) <= self.tracking_length); 25 | } 26 | 27 | /// Add a sample to the tracker. 28 | pub fn add_sample(&mut self) { 29 | self.samples.push(Instant::now()); 30 | self.remove_expired_samples(); 31 | } 32 | 33 | /// Remove the oldest sample from the tracker. 34 | pub fn remove_sample(&mut self) { 35 | if !self.samples.is_empty() { 36 | self.samples.remove(0); 37 | } 38 | self.remove_expired_samples(); 39 | } 40 | 41 | /// Fetches the amount of current non-expired samples. 42 | pub fn get_rate(&mut self) -> usize { 43 | self.remove_expired_samples(); 44 | self.samples.len() 45 | } 46 | } 47 | -------------------------------------------------------------------------------- /assyst-common/src/util/regex.rs: -------------------------------------------------------------------------------- 1 | use lazy_static::lazy_static; 2 | use regex::Regex; 3 | 4 | lazy_static! { 5 | pub static ref CUSTOM_EMOJI: Regex = Regex::new(r"").unwrap(); 6 | pub static ref TENOR_GIF: Regex = Regex::new(r"https://\w+\.tenor\.com/[\w\-]+/[^\.]+\.gif").unwrap(); 7 | pub static ref URL: Regex = Regex::new( 8 | r"https?://(www\.)?[-a-zA-Z0-9@:%._\+~#=]{1,256}\.[a-zA-Z0-9()]{1,6}\b([-a-zA-Z0-9()@:%_\+.~#?&//=]*)" 9 | ) 10 | .unwrap(); 11 | pub static ref USER_MENTION: Regex = Regex::new(r"(?:<@!?)?(\d{16,20})>?").unwrap(); 12 | pub static ref TIME_STRING: Regex = Regex::new("(\\d+)([smhd])").unwrap(); 13 | pub static ref COMMAND_FLAG: Regex = Regex::new(r#"\s+-(\w+)(?: *"([^"]+)"| *([^\-\s]+))?"#).unwrap(); 14 | } 15 | -------------------------------------------------------------------------------- /assyst-common/src/util/table.rs: -------------------------------------------------------------------------------- 1 | /// Returns the longer string of the two given strings 2 | fn get_longer_str<'a>(a: &'a str, b: &'a str) -> &'a str { 3 | if a.len() > b.len() { a } else { b } 4 | } 5 | 6 | /// Generates a table given a list of tuples containing strings 7 | pub fn key_value(input: &[(impl AsRef, impl AsRef)]) -> String { 8 | let longest: &str = input 9 | .iter() 10 | .map(|(x, y)| (x.as_ref(), y.as_ref())) 11 | .fold(input[0].0.as_ref(), |previous, (current, _)| { 12 | get_longer_str(previous, current) 13 | }); 14 | 15 | input 16 | .iter() 17 | .map(|(key, value)| { 18 | format!( 19 | "{}{}: {}\n", 20 | " ".repeat(longest.len() - key.as_ref().len()), 21 | key.as_ref(), 22 | value.as_ref() 23 | ) 24 | }) 25 | .fold(String::new(), |a, b| a + &b) 26 | } 27 | 28 | /// Generates a table given a list of tuples containing strings 29 | pub fn generate_table>(input: &[(T, T)]) -> String { 30 | let longest: &str = input.iter().fold(input[0].0.as_ref(), |previous, (current, _)| { 31 | get_longer_str(previous, current.as_ref()) 32 | }); 33 | 34 | input 35 | .iter() 36 | .map(|(key, value)| { 37 | format!( 38 | "{}{}: {}\n", 39 | " ".repeat(longest.len() - key.as_ref().len()), 40 | key.as_ref(), 41 | value.as_ref() 42 | ) 43 | }) 44 | .fold(String::new(), |a, b| a + &b) 45 | } 46 | 47 | /// Generates a list given a list of tuples containing strings 48 | pub fn generate_list, V: AsRef>(key_name: &str, value_name: &str, values: &[(K, V)]) -> String { 49 | generate_list_fixed_delim(key_name, value_name, values, key_name.len(), value_name.len()) 50 | } 51 | 52 | /// Generates a list given a list of tuples containing strings 53 | pub fn generate_list_fixed_delim, V: AsRef>( 54 | key_name: &str, 55 | value_name: &str, 56 | values: &[(K, V)], 57 | key_delim_len: usize, 58 | value_delim_len: usize, 59 | ) -> String { 60 | let longest = get_longer_str( 61 | key_name, 62 | values.iter().fold(values[0].0.as_ref(), |previous, (current, _)| { 63 | get_longer_str(previous, current.as_ref()) 64 | }), 65 | ); 66 | 67 | let mut output = format!( 68 | " {4}{}\t{}\n {4}{}\t{}", 69 | key_name, 70 | value_name, 71 | "-".repeat(key_delim_len), 72 | "-".repeat(value_delim_len), 73 | " ".repeat(longest.len() - key_name.len()), 74 | ); 75 | 76 | let formatted_values = values 77 | .iter() 78 | .map(|(k, v)| { 79 | format!( 80 | " {}{}\t{}", 81 | " ".repeat(longest.len() - k.as_ref().chars().count()), 82 | k.as_ref(), 83 | v.as_ref() 84 | ) 85 | }) 86 | .collect::>() 87 | .join("\n"); 88 | 89 | output = format!("{output}\n{formatted_values}"); 90 | 91 | output 92 | } 93 | -------------------------------------------------------------------------------- /assyst-core/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "assyst-core" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | anyhow = { workspace = true } 10 | assyst-common = { path = "../assyst-common" } 11 | assyst-database = { path = "../assyst-database" } 12 | assyst-flux-iface = { path = "../assyst-flux-iface" } 13 | assyst-proc-macro = { path = "../assyst-proc-macro" } 14 | assyst-string-fmt = { path = "../assyst-string-fmt" } 15 | assyst-tag = { path = "../assyst-tag" } 16 | assyst-webserver = { path = "../assyst-webserver" } 17 | async-trait = "0.1.77" 18 | aws-lc-rs = "1.10.0" 19 | bincode = "1.3.3" 20 | bytes = "1.5.0" 21 | dash_rt = { git = "https://github.com/y21/dash", rev = "f3fe12b" } 22 | dash_vm = { git = "https://github.com/y21/dash", rev = "f3fe12b" } 23 | emoji = { git = "https://github.com/Jacherr/emoji-rs/", package = "emoji" } 24 | futures-util = "0.3.30" 25 | human_bytes = { version = "0.4", default-features = false } 26 | jemallocator = "0.5.4" 27 | lazy_static = "1.4.0" 28 | libc = "0.2.155" 29 | moka = { version = "0.12.3", features = ["sync"] } 30 | num_cpus = "1.16.0" 31 | paste = "1.0.14" 32 | prometheus = "0.13.3" 33 | rand = "0.8.5" 34 | reqwest = { version = "0.11.24", features = ["json", "stream", "multipart"] } 35 | rustls = "0.23.15" 36 | serde = { workspace = true } 37 | serde_json = "1.0.113" 38 | time = { version = "0.3.31", features = ["macros"] } 39 | tl = "0.7.8" 40 | tokio = { workspace = true } 41 | toml = "0.8.14" 42 | tracing = { workspace = true } 43 | tracing-subscriber = { version = "0.3.16", features = ["time", "env-filter"] } 44 | twilight-gateway = { workspace = true } 45 | twilight-http = { workspace = true } 46 | twilight-model = { workspace = true } 47 | twilight-util = { workspace = true } 48 | url = "2.5.0" 49 | urlencoding = "2.1.3" 50 | zip = "2.1.4" 51 | 52 | [lints] 53 | workspace = true 54 | -------------------------------------------------------------------------------- /assyst-core/src/command/autocomplete.rs: -------------------------------------------------------------------------------- 1 | use twilight_model::id::marker::GuildMarker; 2 | use twilight_model::id::Id; 3 | use twilight_model::user::User; 4 | 5 | pub struct AutocompleteData { 6 | pub guild_id: Option>, 7 | pub user: User, 8 | pub subcommand: Option, 9 | } 10 | 11 | pub const SUGG_LIMIT: usize = 25; 12 | -------------------------------------------------------------------------------- /assyst-core/src/command/flags.rs: -------------------------------------------------------------------------------- 1 | use std::collections::HashMap; 2 | 3 | use anyhow::{bail, Context}; 4 | 5 | #[macro_export] 6 | macro_rules! int_arg_u64 { 7 | ($ctxt:expr, $s:expr, $d:expr) => {{ 8 | let inner = $ctxt 9 | .option_by_name($s) 10 | .map(|o| o.value.clone()) 11 | .unwrap_or(twilight_model::application::interaction::application_command::CommandOptionValue::Integer($d)); 12 | 13 | let inner = 14 | if let twilight_model::application::interaction::application_command::CommandOptionValue::Integer(option) = 15 | inner 16 | { 17 | option as u64 18 | } else { 19 | panic!("download {} wrong arg type", $s); 20 | }; 21 | 22 | inner 23 | }}; 24 | } 25 | 26 | #[macro_export] 27 | macro_rules! int_arg_u64_opt { 28 | ($ctxt:expr, $s:expr) => {{ 29 | let inner = $ctxt.option_by_name($s).map(|o| o.value.clone()); 30 | 31 | let inner = if let Ok( 32 | twilight_model::application::interaction::application_command::CommandOptionValue::Integer(option), 33 | ) = inner 34 | { 35 | Some(option as u64) 36 | } else { 37 | None 38 | }; 39 | 40 | inner 41 | }}; 42 | } 43 | 44 | #[macro_export] 45 | macro_rules! int_arg_bool { 46 | ($ctxt:expr, $s:expr, $d:expr) => {{ 47 | let inner = $ctxt 48 | .option_by_name($s) 49 | .map(|o| o.value.clone()) 50 | .unwrap_or(twilight_model::application::interaction::application_command::CommandOptionValue::Boolean($d)); 51 | 52 | let inner = 53 | if let twilight_model::application::interaction::application_command::CommandOptionValue::Boolean(option) = 54 | inner 55 | { 56 | option 57 | } else { 58 | panic!("download {} wrong arg type", $s); 59 | }; 60 | 61 | inner 62 | }}; 63 | } 64 | 65 | pub enum FlagType { 66 | WithValue, 67 | NoValue, 68 | } 69 | 70 | type ValidFlags = HashMap<&'static str, FlagType>; 71 | 72 | pub trait FlagDecode { 73 | fn from_str(input: &str) -> anyhow::Result 74 | where 75 | Self: Sized; 76 | } 77 | 78 | pub fn flags_from_str(input: &str, valid_flags: ValidFlags) -> anyhow::Result>> { 79 | let args = input.split_ascii_whitespace(); 80 | let mut current_flag: Option = None; 81 | let mut entries: HashMap> = HashMap::new(); 82 | 83 | for arg in args { 84 | if (arg.starts_with("--") && arg.len() > 2) || (arg.starts_with("—") && arg.len() > 1) { 85 | let arglen = if arg.starts_with("--") { 2 } else { 1 }; 86 | 87 | // prev flag present but no value, write to hashmap 88 | if let Some(ref c) = current_flag { 89 | let flag = valid_flags 90 | .get(&c.as_ref()) 91 | .context(format!("Unrecognised flag: {c}"))?; 92 | 93 | if let FlagType::NoValue = flag { 94 | entries.insert(c.clone(), None); 95 | current_flag = Some(arg.chars().skip(arglen).collect::()); 96 | } else { 97 | bail!("Flag {c} expects a value, but none was provided"); 98 | } 99 | } else { 100 | current_flag = Some(arg.chars().skip(arglen).collect::()); 101 | } 102 | } else { 103 | // current flag present, this arg is its value 104 | if let Some(ref c) = current_flag { 105 | let flag = valid_flags 106 | .get(&c.as_ref()) 107 | .context(format!("Unrecognised flag: {c}"))?; 108 | 109 | if let FlagType::WithValue = flag { 110 | entries.insert(c.clone(), Some(arg.to_owned())); 111 | current_flag = None; 112 | } else { 113 | bail!("Flag {c} does not expect a value, even though one was provided"); 114 | } 115 | } 116 | } 117 | } 118 | 119 | // handle case where we assign current flag in last arg, and return 120 | if let Some(c) = current_flag { 121 | let flag = valid_flags 122 | .get(&c.as_ref()) 123 | .context(format!("Unrecognised flag: {c}"))?; 124 | if let FlagType::WithValue = flag { 125 | bail!("Flag {c} expects a value, but none was provided"); 126 | } else { 127 | entries.insert(c.clone(), None); 128 | } 129 | } 130 | 131 | Ok(entries) 132 | } 133 | -------------------------------------------------------------------------------- /assyst-core/src/command/fun/mod.rs: -------------------------------------------------------------------------------- 1 | use std::time::Duration; 2 | 3 | use anyhow::{bail, Context}; 4 | use assyst_proc_macro::command; 5 | 6 | use super::arguments::ImageUrl; 7 | use super::CommandCtxt; 8 | use crate::command::{Availability, Category}; 9 | use crate::rest::audio_identification::identify_song_notsoidentify; 10 | use crate::rest::identify::identify_image; 11 | 12 | pub mod colour; 13 | pub mod translation; 14 | 15 | #[command( 16 | description = "Find a song in a video", 17 | cooldown = Duration::from_secs(2), 18 | access = Availability::Public, 19 | category = Category::Fun, 20 | usage = "[video]", 21 | examples = ["https://link.to.my/video.mp4"], 22 | send_processing = true, 23 | context_menu_message_command = "Find Song" 24 | )] 25 | pub async fn findsong(ctxt: CommandCtxt<'_>, audio: ImageUrl) -> anyhow::Result<()> { 26 | const VALID_FILES: &[&str] = &[".mp3", ".mp4", ".webm", ".ogg", ".wav", ".mov", ".mkv"]; 27 | 28 | if VALID_FILES.iter().all(|x| !audio.0.contains(x)) { 29 | bail!("Finding audio is only supported on audio and video files."); 30 | } 31 | 32 | let result = identify_song_notsoidentify(&ctxt.assyst().reqwest_client, audio.0) 33 | .await 34 | .context("Failed to identify song")?; 35 | 36 | if !result.is_empty() { 37 | let formatted = format!( 38 | "**Title:** {}\n**Artist(s):** {}\n**YouTube Link:** <{}>", 39 | result[0].title.clone(), 40 | result[0] 41 | .artists 42 | .iter() 43 | .map(|x| x.name.clone()) 44 | .collect::>() 45 | .join(", "), 46 | match &result[0].platforms.youtube { 47 | Some(x) => x.url.clone(), 48 | None => "Unknown".to_owned(), 49 | } 50 | ); 51 | ctxt.reply(formatted).await?; 52 | } else { 53 | ctxt.reply("No results found").await?; 54 | } 55 | 56 | Ok(()) 57 | } 58 | 59 | #[command( 60 | description = "AI identify an image", 61 | cooldown = Duration::from_secs(2), 62 | access = Availability::Public, 63 | category = Category::Fun, 64 | usage = "[image]", 65 | examples = ["https://link.to.my/image.png"], 66 | send_processing = true, 67 | context_menu_command = "Identify Image" 68 | )] 69 | pub async fn identify(ctxt: CommandCtxt<'_>, input: ImageUrl) -> anyhow::Result<()> { 70 | let result = identify_image(&ctxt.assyst().reqwest_client, &input.0) 71 | .await 72 | .context("Failed to identify image")?; 73 | 74 | if let Some(d) = result.description { 75 | let formatted = d 76 | .captions 77 | .iter() 78 | .map(|x| format!("I think it's {} ({:.1}% confidence)", x.text, (x.confidence * 100.0))) 79 | .collect::>() 80 | .join("\n"); 81 | 82 | ctxt.reply(formatted).await?; 83 | } else { 84 | ctxt.reply("I really can't describe the picture :flushed:").await?; 85 | } 86 | 87 | Ok(()) 88 | } 89 | -------------------------------------------------------------------------------- /assyst-core/src/command/image/audio.rs: -------------------------------------------------------------------------------- 1 | use std::time::Duration; 2 | 3 | use assyst_proc_macro::command; 4 | 5 | use crate::command::arguments::Image; 6 | use crate::command::{Availability, Category, CommandCtxt}; 7 | 8 | #[command( 9 | description = "give an image drip", 10 | cooldown = Duration::from_secs(3), 11 | access = Availability::Public, 12 | category = Category::Audio, 13 | usage = "[image]", 14 | examples = ["https://link.to.my/image.png"], 15 | send_processing = true 16 | )] 17 | pub async fn drip(ctxt: CommandCtxt<'_>, source: Image) -> anyhow::Result<()> { 18 | let result = ctxt 19 | .flux_handler() 20 | .drip(source.0, ctxt.data.author.id.get(), ctxt.data.guild_id.map(twilight_model::id::Id::get)) 21 | .await?; 22 | 23 | ctxt.reply(result).await?; 24 | 25 | Ok(()) 26 | } 27 | 28 | #[command( 29 | description = "femurbreaker over image", 30 | cooldown = Duration::from_secs(3), 31 | access = Availability::Public, 32 | category = Category::Audio, 33 | usage = "[image]", 34 | examples = ["https://link.to.my/image.png"], 35 | send_processing = true 36 | )] 37 | pub async fn femurbreaker(ctxt: CommandCtxt<'_>, source: Image) -> anyhow::Result<()> { 38 | let result = ctxt 39 | .flux_handler() 40 | .femurbreaker(source.0, ctxt.data.author.id.get(), ctxt.data.guild_id.map(twilight_model::id::Id::get)) 41 | .await?; 42 | 43 | ctxt.reply(result).await?; 44 | 45 | Ok(()) 46 | } 47 | 48 | #[command( 49 | description = "⚠️ alert ⚠️", 50 | cooldown = Duration::from_secs(3), 51 | access = Availability::Public, 52 | category = Category::Audio, 53 | usage = "[image]", 54 | examples = ["https://link.to.my/image.png"], 55 | send_processing = true 56 | )] 57 | pub async fn siren(ctxt: CommandCtxt<'_>, source: Image) -> anyhow::Result<()> { 58 | let result = ctxt 59 | .flux_handler() 60 | .siren(source.0, ctxt.data.author.id.get(), ctxt.data.guild_id.map(twilight_model::id::Id::get)) 61 | .await?; 62 | 63 | ctxt.reply(result).await?; 64 | 65 | Ok(()) 66 | } 67 | 68 | #[command( 69 | description = "give an image some minecraft nostalgia", 70 | cooldown = Duration::from_secs(3), 71 | access = Availability::Public, 72 | category = Category::Audio, 73 | usage = "[image]", 74 | examples = ["https://link.to.my/image.png"], 75 | send_processing = true 76 | )] 77 | pub async fn sweden(ctxt: CommandCtxt<'_>, source: Image) -> anyhow::Result<()> { 78 | let result = ctxt 79 | .flux_handler() 80 | .sweden(source.0, ctxt.data.author.id.get(), ctxt.data.guild_id.map(twilight_model::id::Id::get)) 81 | .await?; 82 | 83 | ctxt.reply(result).await?; 84 | 85 | Ok(()) 86 | } 87 | 88 | #[command( 89 | description = "give your image a grassy theme tune", 90 | cooldown = Duration::from_secs(3), 91 | access = Availability::Public, 92 | category = Category::Audio, 93 | usage = "[image]", 94 | examples = ["https://link.to.my/image.png"], 95 | send_processing = true 96 | )] 97 | pub async fn terraria(ctxt: CommandCtxt<'_>, source: Image) -> anyhow::Result<()> { 98 | let result = ctxt 99 | .flux_handler() 100 | .terraria(source.0, ctxt.data.author.id.get(), ctxt.data.guild_id.map(twilight_model::id::Id::get)) 101 | .await?; 102 | 103 | ctxt.reply(result).await?; 104 | 105 | Ok(()) 106 | } 107 | -------------------------------------------------------------------------------- /assyst-core/src/command/image/bloom.rs: -------------------------------------------------------------------------------- 1 | use std::collections::HashMap; 2 | use std::time::Duration; 3 | 4 | use anyhow::Context; 5 | use assyst_proc_macro::command; 6 | use twilight_util::builder::command::IntegerBuilder; 7 | 8 | use crate::command::arguments::{Image, ParseArgument}; 9 | use crate::command::errors::TagParseError; 10 | use crate::command::flags::{flags_from_str, FlagDecode, FlagType}; 11 | use crate::command::{Availability, Category, CommandCtxt}; 12 | use crate::int_arg_u64_opt; 13 | 14 | #[derive(Default)] 15 | pub struct BloomFlags { 16 | pub radius: Option, 17 | pub brightness: Option, 18 | pub sharpness: Option, 19 | } 20 | impl FlagDecode for BloomFlags { 21 | fn from_str(input: &str) -> anyhow::Result { 22 | let mut valid_flags = HashMap::new(); 23 | valid_flags.insert("radius", FlagType::WithValue); 24 | valid_flags.insert("sharpness", FlagType::WithValue); 25 | valid_flags.insert("brightness", FlagType::WithValue); 26 | 27 | let raw_decode = flags_from_str(input, valid_flags)?; 28 | let result = Self { 29 | radius: raw_decode 30 | .get("radius") 31 | .unwrap_or(&None) 32 | .clone() 33 | .map(|x| x.parse().context("Provided radius is invalid")) 34 | .transpose()?, 35 | sharpness: raw_decode 36 | .get("sharpness") 37 | .unwrap_or(&None) 38 | .clone() 39 | .map(|x| x.parse().context("Provided sharpness is invalid")) 40 | .transpose()?, 41 | brightness: raw_decode 42 | .get("brightness") 43 | .unwrap_or(&None) 44 | .clone() 45 | .map(|x| x.parse().context("Provided brightness is invalid")) 46 | .transpose()?, 47 | }; 48 | 49 | Ok(result) 50 | } 51 | } 52 | impl ParseArgument for BloomFlags { 53 | fn as_command_options(_: &str) -> Vec { 54 | vec![ 55 | IntegerBuilder::new("radius", "bloom radius").required(false).build(), 56 | IntegerBuilder::new("sharpness", "bloom sharpness") 57 | .required(false) 58 | .build(), 59 | IntegerBuilder::new("brightness", "bloom brightness") 60 | .required(false) 61 | .build(), 62 | ] 63 | } 64 | 65 | async fn parse_raw_message( 66 | ctxt: &mut crate::command::RawMessageParseCtxt<'_>, 67 | label: crate::command::Label, 68 | ) -> Result { 69 | let args = ctxt.rest_all(label); 70 | let parsed = Self::from_str(&args).map_err(TagParseError::FlagParseError)?; 71 | Ok(parsed) 72 | } 73 | 74 | async fn parse_command_option( 75 | ctxt: &mut crate::command::InteractionCommandParseCtxt<'_>, 76 | _: crate::command::Label, 77 | ) -> Result { 78 | let radius = int_arg_u64_opt!(ctxt, "radius"); 79 | let sharpness = int_arg_u64_opt!(ctxt, "sharpness"); 80 | let brightness = int_arg_u64_opt!(ctxt, "brightness"); 81 | 82 | Ok(Self { radius, brightness, sharpness }) 83 | } 84 | } 85 | 86 | #[command( 87 | description = "add bloom to an image", 88 | cooldown = Duration::from_secs(2), 89 | access = Availability::Public, 90 | category = Category::Image, 91 | usage = "[image] ", 92 | examples = ["https://link.to.my/image.png", "https://link.to.my/image.png --brightness 100 --sharpness 25 --radius 10"], 93 | send_processing = true, 94 | flag_descriptions = [ 95 | ("radius", "Bloom radius as a number"), 96 | ("brightness", "Bloom brightness as a number"), 97 | ("sharpness", "Bloom sharpness as a number"), 98 | ] 99 | )] 100 | pub async fn bloom(ctxt: CommandCtxt<'_>, source: Image, flags: BloomFlags) -> anyhow::Result<()> { 101 | let result = ctxt 102 | .flux_handler() 103 | .bloom( 104 | source.0, 105 | flags.radius, 106 | flags.sharpness, 107 | flags.brightness, 108 | ctxt.data.author.id.get(), 109 | ctxt.data.guild_id.map(twilight_model::id::Id::get), 110 | ) 111 | .await?; 112 | 113 | ctxt.reply(result).await?; 114 | 115 | Ok(()) 116 | } 117 | -------------------------------------------------------------------------------- /assyst-core/src/command/image/caption.rs: -------------------------------------------------------------------------------- 1 | use std::collections::HashMap; 2 | use std::time::Duration; 3 | 4 | use assyst_proc_macro::command; 5 | use twilight_util::builder::command::BooleanBuilder; 6 | 7 | use crate::command::arguments::{Image, ParseArgument, Rest}; 8 | use crate::command::errors::TagParseError; 9 | use crate::command::flags::{flags_from_str, FlagDecode, FlagType}; 10 | use crate::command::{Availability, Category, CommandCtxt}; 11 | use crate::int_arg_bool; 12 | 13 | #[command( 14 | description = "add a caption to an image", 15 | cooldown = Duration::from_secs(2), 16 | access = Availability::Public, 17 | category = Category::Image, 18 | usage = "[image] [caption] <...flags>", 19 | examples = ["https://link.to.my/image.png hello there", "https://link.to.my/image.png i am on the bottom --bottom", "https://link.to.my/image.png i am an inverted caption --black"], 20 | send_processing = true, 21 | flag_descriptions = [ 22 | ("bottom", "Setting this flag puts the caption on the bottom of the image"), 23 | ("black", "Setting this flag inverts the caption"), 24 | ] 25 | )] 26 | pub async fn caption(ctxt: CommandCtxt<'_>, source: Image, text: Rest, flags: CaptionFlags) -> anyhow::Result<()> { 27 | let result = ctxt 28 | .flux_handler() 29 | .caption( 30 | source.0, 31 | text.0, 32 | flags.bottom, 33 | flags.black, 34 | ctxt.data.author.id.get(), 35 | ctxt.data.guild_id.map(twilight_model::id::Id::get), 36 | ) 37 | .await?; 38 | 39 | ctxt.reply(result).await?; 40 | 41 | Ok(()) 42 | } 43 | 44 | #[derive(Default)] 45 | pub struct CaptionFlags { 46 | pub bottom: bool, 47 | pub black: bool, 48 | } 49 | impl FlagDecode for CaptionFlags { 50 | fn from_str(input: &str) -> anyhow::Result 51 | where 52 | Self: Sized, 53 | { 54 | let mut valid_flags = HashMap::new(); 55 | valid_flags.insert("bottom", FlagType::NoValue); 56 | valid_flags.insert("black", FlagType::NoValue); 57 | 58 | let raw_decode = flags_from_str(input, valid_flags)?; 59 | 60 | let result = Self { 61 | bottom: raw_decode.contains_key("bottom"), 62 | black: raw_decode.contains_key("black"), 63 | }; 64 | 65 | Ok(result) 66 | } 67 | } 68 | impl ParseArgument for CaptionFlags { 69 | fn as_command_options(_: &str) -> Vec { 70 | vec![ 71 | BooleanBuilder::new("bottom", "put the caption on the bottom") 72 | .required(false) 73 | .build(), 74 | BooleanBuilder::new("black", "invert the caption") 75 | .required(false) 76 | .build(), 77 | ] 78 | } 79 | 80 | async fn parse_raw_message( 81 | ctxt: &mut crate::command::RawMessageParseCtxt<'_>, 82 | label: crate::command::Label, 83 | ) -> Result { 84 | let args = ctxt.rest_all(label); 85 | let parsed = Self::from_str(&args).map_err(TagParseError::FlagParseError)?; 86 | Ok(parsed) 87 | } 88 | 89 | async fn parse_command_option( 90 | ctxt: &mut crate::command::InteractionCommandParseCtxt<'_>, 91 | _: crate::command::Label, 92 | ) -> Result { 93 | let bottom = int_arg_bool!(ctxt, "bottom", false); 94 | let black = int_arg_bool!(ctxt, "black", false); 95 | 96 | Ok(Self { bottom, black }) 97 | } 98 | } 99 | -------------------------------------------------------------------------------- /assyst-core/src/command/image/randomize.rs: -------------------------------------------------------------------------------- 1 | use std::collections::HashMap; 2 | use std::time::Duration; 3 | 4 | use anyhow::Context; 5 | use assyst_flux_iface::flux_request::FluxRequest; 6 | use assyst_proc_macro::command; 7 | use rand::{thread_rng, Rng}; 8 | 9 | use crate::command::arguments::Image; 10 | use crate::command::{Availability, Category, CommandCtxt}; 11 | 12 | const VALID_EFFECTS: &[&str] = &[ 13 | "bloom", 14 | "blur", 15 | "deepfry", 16 | "fisheye", 17 | "flip", 18 | "flop", 19 | "globe", 20 | "grayscale", 21 | "invert", 22 | "jpeg", 23 | "magik", 24 | "neon", 25 | "paint", 26 | "pixelate", 27 | "rainbow", 28 | ]; 29 | 30 | #[command( 31 | description = "apply random effects to an image", 32 | aliases = ["random", "randomise", "badcmd"], 33 | cooldown = Duration::from_secs(2), 34 | access = Availability::Public, 35 | category = Category::Image, 36 | usage = "[image] ", 37 | examples = ["https://link.to.my/image.png 3"], 38 | send_processing = true 39 | )] 40 | pub async fn randomize(ctxt: CommandCtxt<'_>, source: Image, count: Option) -> anyhow::Result<()> { 41 | let mut effects: Vec<&str> = Vec::new(); 42 | 43 | for _ in 0..count.unwrap_or(3).clamp(1, 5) { 44 | let next = loop { 45 | let tmp = VALID_EFFECTS[thread_rng().gen_range(0..VALID_EFFECTS.len())]; 46 | if effects.last() != Some(&tmp) { 47 | break tmp; 48 | } 49 | }; 50 | 51 | effects.push(next); 52 | } 53 | 54 | let limits = ctxt 55 | .assyst() 56 | .flux_handler 57 | .get_request_limits(ctxt.data.author.id.get(), ctxt.data.guild_id.map(twilight_model::id::Id::get)) 58 | .await?; 59 | 60 | let mut request = FluxRequest::new_with_input_and_limits(source.0, &limits); 61 | for e in &effects { 62 | request.operation((*e).to_string(), HashMap::new()); 63 | } 64 | 65 | request.output(); 66 | 67 | let result = ctxt 68 | .assyst() 69 | .flux_handler 70 | .run_flux(request, limits.time) 71 | .await 72 | .context(format!("Applied effects: {}", effects.join(", ")))?; 73 | 74 | ctxt.reply(( 75 | result, 76 | &format!( 77 | "Applied effects: {}", 78 | effects.iter().map(|e| format!("`{e}`")).collect::>().join(", ") 79 | )[..], 80 | )) 81 | .await?; 82 | 83 | Ok(()) 84 | } 85 | -------------------------------------------------------------------------------- /assyst-core/src/command/image/speechbubble.rs: -------------------------------------------------------------------------------- 1 | use std::collections::HashMap; 2 | use std::time::Duration; 3 | 4 | use assyst_proc_macro::command; 5 | use twilight_util::builder::command::BooleanBuilder; 6 | 7 | use crate::command::arguments::{Image, ParseArgument}; 8 | use crate::command::errors::TagParseError; 9 | use crate::command::flags::{flags_from_str, FlagDecode, FlagType}; 10 | use crate::command::{Availability, Category, CommandCtxt}; 11 | use crate::int_arg_bool; 12 | 13 | #[command( 14 | description = "add a speechbubble to an image", 15 | aliases = ["speech", "bubble"], 16 | cooldown = Duration::from_secs(2), 17 | access = Availability::Public, 18 | category = Category::Image, 19 | usage = "[image] <...flags>", 20 | examples = ["https://link.to.my/image.png", "https://link.to.my/image.png --solid"], 21 | send_processing = true, 22 | flag_descriptions = [ 23 | ("solid", "Setting this flag will make the speech bubble a solid white instead of transparent"), 24 | ] 25 | )] 26 | pub async fn speechbubble(ctxt: CommandCtxt<'_>, source: Image, flags: SpeechBubbleFlags) -> anyhow::Result<()> { 27 | let result = ctxt 28 | .flux_handler() 29 | .speech_bubble( 30 | source.0, 31 | flags.solid, 32 | ctxt.data.author.id.get(), 33 | ctxt.data.guild_id.map(twilight_model::id::Id::get), 34 | ) 35 | .await?; 36 | 37 | ctxt.reply(result).await?; 38 | 39 | Ok(()) 40 | } 41 | 42 | #[derive(Default)] 43 | pub struct SpeechBubbleFlags { 44 | pub solid: bool, 45 | } 46 | impl FlagDecode for SpeechBubbleFlags { 47 | fn from_str(input: &str) -> anyhow::Result 48 | where 49 | Self: Sized, 50 | { 51 | let mut valid_flags = HashMap::new(); 52 | valid_flags.insert("solid", FlagType::NoValue); 53 | 54 | let raw_decode = flags_from_str(input, valid_flags)?; 55 | 56 | let result = Self { 57 | solid: raw_decode.contains_key("solid"), 58 | }; 59 | 60 | Ok(result) 61 | } 62 | } 63 | impl ParseArgument for SpeechBubbleFlags { 64 | fn as_command_options(_: &str) -> Vec { 65 | vec![ 66 | BooleanBuilder::new("solid", "make the speech bubble solid") 67 | .required(false) 68 | .build(), 69 | ] 70 | } 71 | 72 | async fn parse_raw_message( 73 | ctxt: &mut crate::command::RawMessageParseCtxt<'_>, 74 | label: crate::command::Label, 75 | ) -> Result { 76 | let args = ctxt.rest_all(label); 77 | let parsed = Self::from_str(&args).map_err(TagParseError::FlagParseError)?; 78 | Ok(parsed) 79 | } 80 | 81 | async fn parse_command_option( 82 | ctxt: &mut crate::command::InteractionCommandParseCtxt<'_>, 83 | _: crate::command::Label, 84 | ) -> Result { 85 | let solid = int_arg_bool!(ctxt, "solid", false); 86 | 87 | Ok(Self { solid }) 88 | } 89 | } 90 | -------------------------------------------------------------------------------- /assyst-core/src/command/messagebuilder.rs: -------------------------------------------------------------------------------- 1 | use assyst_common::util::filetype::{get_sig, Type}; 2 | use twilight_model::channel::message::Component; 3 | 4 | use super::arguments::Image; 5 | use super::componentctxt::ComponentCtxtRegister; 6 | 7 | #[derive(Debug)] 8 | pub struct Attachment { 9 | pub name: Box, 10 | pub data: Vec, 11 | } 12 | 13 | impl From for Attachment { 14 | fn from(value: Image) -> Self { 15 | let ext = get_sig(&value.0).unwrap_or(Type::PNG).as_str(); 16 | Attachment { 17 | name: format!("attachment.{ext}").into(), 18 | data: value.0, 19 | } 20 | } 21 | } 22 | 23 | pub struct MessageBuilder { 24 | pub content: Option, 25 | pub attachment: Option, 26 | pub components: Option>, 27 | pub component_ctxt: Option, 28 | } 29 | 30 | impl From<&str> for MessageBuilder { 31 | fn from(value: &str) -> Self { 32 | Self { 33 | content: Some(value.into()), 34 | attachment: None, 35 | components: None, 36 | component_ctxt: None, 37 | } 38 | } 39 | } 40 | impl From for MessageBuilder { 41 | fn from(value: String) -> Self { 42 | Self { 43 | content: Some(value), 44 | attachment: None, 45 | components: None, 46 | component_ctxt: None, 47 | } 48 | } 49 | } 50 | 51 | impl From for MessageBuilder { 52 | fn from(value: Attachment) -> Self { 53 | Self { 54 | content: None, 55 | attachment: Some(value), 56 | components: None, 57 | component_ctxt: None, 58 | } 59 | } 60 | } 61 | 62 | impl From<(Attachment, String)> for MessageBuilder { 63 | fn from(value: (Attachment, String)) -> Self { 64 | Self { 65 | content: Some(value.1), 66 | attachment: Some(value.0), 67 | components: None, 68 | component_ctxt: None, 69 | } 70 | } 71 | } 72 | 73 | impl From for MessageBuilder { 74 | fn from(value: Image) -> Self { 75 | Self { 76 | content: None, 77 | attachment: Some(value.into()), 78 | components: None, 79 | component_ctxt: None, 80 | } 81 | } 82 | } 83 | impl From<(Image, &str)> for MessageBuilder { 84 | fn from((image, text): (Image, &str)) -> Self { 85 | Self { 86 | attachment: Some(image.into()), 87 | content: Some(text.into()), 88 | components: None, 89 | component_ctxt: None, 90 | } 91 | } 92 | } 93 | impl From> for MessageBuilder { 94 | fn from(value: Vec) -> Self { 95 | Self { 96 | attachment: Some(Image(value).into()), 97 | content: None, 98 | components: None, 99 | component_ctxt: None, 100 | } 101 | } 102 | } 103 | impl From<(Vec, &str)> for MessageBuilder { 104 | fn from((value, text): (Vec, &str)) -> Self { 105 | Self { 106 | attachment: Some(Image(value).into()), 107 | content: Some(text.into()), 108 | components: None, 109 | component_ctxt: None, 110 | } 111 | } 112 | } 113 | -------------------------------------------------------------------------------- /assyst-core/src/command/misc/btchannel.rs: -------------------------------------------------------------------------------- 1 | use std::time::Duration; 2 | 3 | use anyhow::{bail, ensure, Context}; 4 | use assyst_common::util::discord::ensure_same_guild; 5 | use assyst_database::model::badtranslator_channel::BadTranslatorChannel; 6 | use assyst_proc_macro::command; 7 | 8 | use crate::command::arguments::{Channel, Word}; 9 | use crate::command::{Availability, Category, CommandCtxt}; 10 | use crate::define_commandgroup; 11 | use crate::rest::bad_translation::validate_language; 12 | 13 | #[command( 14 | description = "register a new badtranslator channel", 15 | aliases = ["create"], 16 | cooldown = Duration::from_secs(10), 17 | access = Availability::ServerManagers, 18 | category = Category::Misc, 19 | usage = "[channel] [language]", 20 | examples = ["#bt en"], 21 | guild_only = true 22 | )] 23 | pub async fn add(ctxt: CommandCtxt<'_>, channel: Channel, target_language: Word) -> anyhow::Result<()> { 24 | let Some(guild_id) = ctxt.data.guild_id else { 25 | bail!("BadTranslator channels can ony be created inside of servers."); 26 | }; 27 | 28 | ensure_same_guild(&ctxt.assyst().http_client, channel.0.id.get(), guild_id.get()).await?; 29 | ensure!( 30 | validate_language(&ctxt.assyst().reqwest_client, &target_language.0).await?, 31 | "This language does not exist or cannot be used as a target language. Run `{}btchannel languages` for a list of languages", 32 | ctxt.data.calling_prefix 33 | ); 34 | 35 | ctxt.assyst() 36 | .http_client 37 | .create_webhook(channel.0.id, "Bad Translator") 38 | .await 39 | .context("Failed to create BadTranslator webhook")?; 40 | 41 | let new = BadTranslatorChannel { 42 | id: channel.0.id.get() as i64, 43 | target_language: target_language.0, 44 | }; 45 | 46 | ensure!( 47 | new.set(&ctxt.assyst().database_handler) 48 | .await 49 | .context("Failed to register BadTranslator channel")?, 50 | "This channel is already registered as a BadTranslator channel." 51 | ); 52 | 53 | ctxt.assyst() 54 | .bad_translator 55 | .add_channel(channel.0.id.get(), &new.target_language) 56 | .await; 57 | 58 | ctxt.reply(format!( 59 | "The channel <#{}> has been registered as a BadTranslator channel.", 60 | channel.0.id.get() 61 | )) 62 | .await?; 63 | 64 | Ok(()) 65 | } 66 | 67 | #[command( 68 | description = "delete an existing badtranslator channel", 69 | aliases = ["remove"], 70 | cooldown = Duration::from_secs(10), 71 | access = Availability::ServerManagers, 72 | category = Category::Misc, 73 | usage = "[channel]", 74 | examples = ["#bt"], 75 | guild_only = true 76 | )] 77 | pub async fn remove(ctxt: CommandCtxt<'_>, channel: Channel) -> anyhow::Result<()> { 78 | let Some(guild_id) = ctxt.data.guild_id else { 79 | bail!("BadTranslator channels can ony be deleted inside of servers."); 80 | }; 81 | 82 | ensure_same_guild(&ctxt.assyst().http_client, channel.0.id.get(), guild_id.get()).await?; 83 | 84 | let webhooks = ctxt 85 | .assyst() 86 | .http_client 87 | .channel_webhooks(channel.0.id) 88 | .await 89 | .context("Failed to get webhooks")? 90 | .model() 91 | .await 92 | .context("Failed to get webhooks")?; 93 | 94 | for webhook in webhooks { 95 | if webhook.name == Some("Bad Translator".to_owned()) { 96 | ctxt.assyst() 97 | .http_client 98 | .delete_webhook(webhook.id) 99 | .await 100 | .context("Failed to delete BadTranslator webhook")?; 101 | } 102 | } 103 | 104 | ensure!( 105 | BadTranslatorChannel::delete(&ctxt.assyst().database_handler, channel.0.id.get() as i64) 106 | .await 107 | .context("Failed to delete BadTranslator channel")?, 108 | "Failed to delete this BadTranslator channel - does it exist?" 109 | ); 110 | 111 | ctxt.assyst().bad_translator.remove_bt_channel(channel.0.id.get()).await; 112 | 113 | ctxt.reply(format!( 114 | "The channel <#{}> has been unregistered as a BadTranslator channel.", 115 | channel.0.id.get() 116 | )) 117 | .await?; 118 | 119 | Ok(()) 120 | } 121 | 122 | define_commandgroup! { 123 | name: btchannel, 124 | access: Availability::ServerManagers, 125 | category: Category::Misc, 126 | description: "manage badtranslator channels", 127 | usage: "[subcommand] ", 128 | guild_only: true, 129 | commands: [ 130 | "add" => add, 131 | "remove" => remove 132 | ] 133 | } 134 | -------------------------------------------------------------------------------- /assyst-core/src/command/misc/prefix.rs: -------------------------------------------------------------------------------- 1 | use std::time::Duration; 2 | 3 | use anyhow::{bail, ensure, Context}; 4 | use assyst_database::model::prefix::Prefix; 5 | use assyst_proc_macro::command; 6 | use assyst_string_fmt::Markdown; 7 | 8 | use crate::command::arguments::Word; 9 | use crate::command::{Availability, Category, CommandCtxt}; 10 | use crate::define_commandgroup; 11 | 12 | #[command( 13 | description = "get server prefix", 14 | access = Availability::Public, 15 | cooldown = Duration::from_secs(2), 16 | category = Category::Misc, 17 | examples = [""], 18 | )] 19 | pub async fn default(ctxt: CommandCtxt<'_>) -> anyhow::Result<()> { 20 | let Some(guild_id) = ctxt.data.guild_id else { 21 | bail!("prefix getting and setting can only be used in guilds") 22 | }; 23 | 24 | let prefix = Prefix::get(&ctxt.assyst().database_handler, guild_id.get()) 25 | .await 26 | .context("Failed to get guild prefix")? 27 | .context("This guild has no set prefix?")?; 28 | 29 | ctxt.reply(format!("This server's prefix is: {}", prefix.prefix.codestring())) 30 | .await?; 31 | 32 | Ok(()) 33 | } 34 | 35 | #[command( 36 | description = "set server prefix", 37 | access = Availability::ServerManagers, 38 | cooldown = Duration::from_secs(2), 39 | category = Category::Misc, 40 | examples = ["-", "%"], 41 | )] 42 | pub async fn set(ctxt: CommandCtxt<'_>, new: Word) -> anyhow::Result<()> { 43 | let Some(guild_id) = ctxt.data.guild_id else { 44 | bail!("Prefix getting and setting can only be used in guilds.") 45 | }; 46 | 47 | ensure!(new.0.len() < 14, "Prefixes cannot be longer than 14 characters."); 48 | 49 | let new = Prefix { prefix: new.0 }; 50 | new.set(&ctxt.assyst().database_handler, guild_id.get()) 51 | .await 52 | .context("Failed to set new prefix")?; 53 | 54 | ctxt.reply(format!("This server's prefix is now: {}", new.prefix.codestring())) 55 | .await?; 56 | 57 | Ok(()) 58 | } 59 | 60 | define_commandgroup! { 61 | name: prefix, 62 | access: Availability::Public, 63 | category: Category::Misc, 64 | description: "get or set server prefix", 65 | usage: " ", 66 | commands: [ 67 | "set" => set 68 | ], 69 | default_interaction_subcommand: "get", 70 | default: default 71 | } 72 | -------------------------------------------------------------------------------- /assyst-core/src/command/misc/remind.rs: -------------------------------------------------------------------------------- 1 | use std::time::{Duration, SystemTime, UNIX_EPOCH}; 2 | 3 | use anyhow::{bail, Context}; 4 | use assyst_common::util::discord::format_discord_timestamp; 5 | use assyst_common::util::format_time; 6 | use assyst_database::model::reminder::Reminder; 7 | use assyst_proc_macro::command; 8 | 9 | use crate::command::arguments::{Rest, Time}; 10 | use crate::command::{Availability, Category, CommandCtxt}; 11 | use crate::define_commandgroup; 12 | 13 | #[command( 14 | aliases = ["reminder"], 15 | description = "set a new reminder - time format is xdyhzm (check examples)", 16 | access = Availability::Public, 17 | cooldown = Duration::from_secs(2), 18 | category = Category::Misc, 19 | usage = "[time] ", 20 | examples = ["2h do the laundry", "3d30m hand assignment in", "30m"], 21 | )] 22 | pub async fn default(ctxt: CommandCtxt<'_>, when: Time, text: Option) -> anyhow::Result<()> { 23 | if when.millis < 1000 { 24 | bail!( 25 | "Invalid time provided (see {}help remind for examples)", 26 | ctxt.data.calling_prefix 27 | ); 28 | } else if when.millis / 1000 / 60 / 60 / 24 / 365 /* years */ >= 100 { 29 | bail!("Cannot set a reminder further than 100 years in the future :-("); 30 | } 31 | 32 | let timestamp = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_millis() as u64 + when.millis; 33 | 34 | let text = text.map_or("...".to_owned(), |x| x.0); 35 | 36 | if text.len() > 250 { 37 | bail!("Reminder message cannot exceed 250 characters."); 38 | } 39 | 40 | let reminder = Reminder { 41 | id: 0, // unused 42 | user_id: ctxt.data.author.id.get() as i64, 43 | timestamp: timestamp as i64, 44 | guild_id: ctxt.data.guild_id.map_or(0, twilight_model::id::Id::get) as i64, 45 | channel_id: ctxt.data.channel_id.get() as i64, 46 | message_id: ctxt.data.message.map_or(0, |x| x.id.get()) as i64, 47 | message: text, 48 | }; 49 | 50 | reminder 51 | .insert(&ctxt.assyst().database_handler) 52 | .await 53 | .context("Failed to insert reminder to database")?; 54 | 55 | ctxt.reply(format!( 56 | "Reminder successfully set for {} from now.", 57 | format_time(when.millis) 58 | )) 59 | .await?; 60 | 61 | Ok(()) 62 | } 63 | 64 | #[command( 65 | description = "list all of your reminders", 66 | access = Availability::Public, 67 | cooldown = Duration::from_secs(2), 68 | category = Category::Misc, 69 | usage = "", 70 | examples = [""], 71 | )] 72 | pub async fn list(ctxt: CommandCtxt<'_>) -> anyhow::Result<()> { 73 | let reminders = Reminder::fetch_user_reminders(&ctxt.assyst().database_handler, ctxt.data.author.id.get(), 10) 74 | .await 75 | .context("Failed to fetch reminders")?; 76 | 77 | if reminders.is_empty() { 78 | ctxt.reply("You don't have any set reminders.").await?; 79 | return Ok(()); 80 | } 81 | 82 | let formatted = reminders.iter().fold(String::new(), |mut f, reminder| { 83 | use std::fmt::Write; 84 | writeln!( 85 | f, 86 | "[#{}] {}: `{}`", 87 | reminder.id, 88 | format_discord_timestamp(reminder.timestamp as u64), 89 | reminder.message 90 | ) 91 | .unwrap(); 92 | f 93 | }); 94 | 95 | ctxt.reply(format!(":calendar: **Upcoming Reminders:**\n\n{formatted}")) 96 | .await?; 97 | 98 | Ok(()) 99 | } 100 | 101 | define_commandgroup! { 102 | name: remind, 103 | access: Availability::Public, 104 | category: Category::Misc, 105 | aliases: ["t"], 106 | cooldown: Duration::from_secs(2), 107 | description: "assyst reminders - get and set reminders", 108 | examples: ["2h do the laundry", "3d30m hand assignment in", "30m"], 109 | usage: "[time] ", 110 | commands: [ 111 | "list" => list 112 | ], 113 | default_interaction_subcommand: "create", 114 | default: default 115 | } 116 | -------------------------------------------------------------------------------- /assyst-core/src/command/services/cooltext.rs: -------------------------------------------------------------------------------- 1 | use std::time::Duration; 2 | 3 | use assyst_proc_macro::command; 4 | use assyst_string_fmt::Markdown; 5 | use rand::{thread_rng, Rng}; 6 | 7 | use crate::assyst::ThreadSafeAssyst; 8 | use crate::command::arguments::{Rest, WordAutocomplete}; 9 | use crate::command::autocomplete::AutocompleteData; 10 | use crate::command::{Availability, Category, CommandCtxt}; 11 | use crate::define_commandgroup; 12 | use crate::rest::cooltext::STYLES; 13 | 14 | pub async fn cooltext_options_autocomplete(_a: ThreadSafeAssyst, _d: AutocompleteData) -> Vec { 15 | let options = STYLES.iter().map(|x| x.0.to_owned()).collect::>(); 16 | options 17 | } 18 | 19 | #[command( 20 | description = "make some cool text", 21 | access = Availability::Public, 22 | cooldown = Duration::from_secs(2), 23 | category = Category::Services, 24 | examples = ["burning hello", "saint fancy", "random im random"], 25 | send_processing = true 26 | )] 27 | pub async fn default( 28 | ctxt: CommandCtxt<'_>, 29 | #[autocomplete = "crate::command::services::cooltext::cooltext_options_autocomplete"] style: WordAutocomplete, 30 | text: Rest, 31 | ) -> anyhow::Result<()> { 32 | let style = if &style.0 == "random" { 33 | let rand = thread_rng().gen_range(0..STYLES.len()); 34 | STYLES[rand].0 35 | } else { 36 | &style.0 37 | }; 38 | 39 | let result = crate::rest::cooltext::cooltext(style, text.0.as_str()).await?; 40 | ctxt.reply((result, &format!("**Style:** `{style}`")[..])).await?; 41 | 42 | Ok(()) 43 | } 44 | 45 | #[command( 46 | description = "list all cooltext options", 47 | access = Availability::Public, 48 | cooldown = Duration::from_secs(2), 49 | category = Category::Services, 50 | usage = "", 51 | examples = [""], 52 | )] 53 | pub async fn list(ctxt: CommandCtxt<'_>) -> anyhow::Result<()> { 54 | let options = STYLES.iter().map(|x| x.0.to_owned()).collect::>(); 55 | 56 | ctxt.reply(format!( 57 | "**All Cooltext supported fonts:**\n{}", 58 | options.join(", ").codeblock("") 59 | )) 60 | .await?; 61 | 62 | Ok(()) 63 | } 64 | 65 | define_commandgroup! { 66 | name: cooltext, 67 | access: Availability::Public, 68 | category: Category::Services, 69 | aliases: ["ct", "funtext"], 70 | cooldown: Duration::from_secs(5), 71 | description: "Write some cool text", 72 | examples: ["random hello", "warp warpy text", "list"], 73 | usage: "[colour]", 74 | commands: [ 75 | "list" => list 76 | ], 77 | default_interaction_subcommand: "create", 78 | default: default 79 | } 80 | -------------------------------------------------------------------------------- /assyst-core/src/command/services/mod.rs: -------------------------------------------------------------------------------- 1 | use std::time::Duration; 2 | 3 | use assyst_proc_macro::command; 4 | 5 | use super::arguments::Rest; 6 | use super::CommandCtxt; 7 | use crate::command::{Availability, Category}; 8 | use crate::rest::cooltext::burn_text; 9 | use crate::rest::r34::get_random_r34; 10 | 11 | pub mod cooltext; 12 | pub mod download; 13 | 14 | #[command( 15 | aliases = ["firetext"], 16 | description = "make some burning text", 17 | access = Availability::Public, 18 | cooldown = Duration::from_secs(2), 19 | category = Category::Services, 20 | usage = "[text]", 21 | examples = ["yep im burning"], 22 | send_processing = true, 23 | context_menu_message_command = "Burn Text" 24 | )] 25 | pub async fn burntext(ctxt: CommandCtxt<'_>, text: Rest) -> anyhow::Result<()> { 26 | let result = burn_text(&text.0).await?; 27 | 28 | ctxt.reply(result).await?; 29 | 30 | Ok(()) 31 | } 32 | 33 | #[command( 34 | name = "rule34", 35 | aliases = ["r34"], 36 | description = "get random image from r34", 37 | access = Availability::Public, 38 | cooldown = Duration::from_secs(2), 39 | category = Category::Services, 40 | usage = "", 41 | examples = ["", "assyst"], 42 | age_restricted = true 43 | )] 44 | pub async fn r34(ctxt: CommandCtxt<'_>, tags: Option) -> anyhow::Result<()> { 45 | let result = get_random_r34(&ctxt.assyst().reqwest_client, &tags.unwrap_or(Rest(String::new())).0).await?; 46 | let reply = format!("{} (Score: **{}**)", result.file_url, result.score); 47 | 48 | ctxt.reply(reply).await?; 49 | 50 | Ok(()) 51 | } 52 | -------------------------------------------------------------------------------- /assyst-core/src/command/source.rs: -------------------------------------------------------------------------------- 1 | #[derive(Clone, Eq, PartialEq)] 2 | pub enum Source { 3 | RawMessage, 4 | Interaction, 5 | } 6 | -------------------------------------------------------------------------------- /assyst-core/src/command_ratelimits.rs: -------------------------------------------------------------------------------- 1 | use std::time::{Duration, Instant}; 2 | 3 | use moka::sync::Cache; 4 | 5 | /// All command ratelimits, in the format <(guild/user id, command name) => time command was 6 | /// ran> 7 | pub struct CommandRatelimits(Cache<(u64, &'static str), Instant>); 8 | impl CommandRatelimits { 9 | pub fn new() -> Self { 10 | Self( 11 | Cache::builder() 12 | .max_capacity(1000) 13 | .time_to_idle(Duration::from_secs(60 * 5)) 14 | .build(), 15 | ) 16 | } 17 | 18 | pub fn insert(&self, id: u64, command_name: &'static str, value: Instant) { 19 | self.0.insert((id, command_name), value); 20 | } 21 | 22 | pub fn get(&self, id: u64, command_name: &'static str) -> Option { 23 | self.0.get(&(id, command_name)) 24 | } 25 | } 26 | -------------------------------------------------------------------------------- /assyst-core/src/downloader.rs: -------------------------------------------------------------------------------- 1 | use core::fmt; 2 | use std::sync::atomic::{AtomicUsize, Ordering}; 3 | use std::time::Duration; 4 | 5 | use assyst_common::config::CONFIG; 6 | use bytes::Bytes; 7 | use futures_util::{Stream, StreamExt}; 8 | use human_bytes::human_bytes; 9 | use reqwest::{Client, StatusCode, Url}; 10 | 11 | pub const ABSOLUTE_INPUT_FILE_SIZE_LIMIT_BYTES: usize = 250_000_000; 12 | static PROXY_NUM: AtomicUsize = AtomicUsize::new(0); 13 | 14 | #[derive(Debug)] 15 | pub enum DownloadError { 16 | ProxyNetworkError, 17 | InvalidStatus, 18 | Url(url::ParseError), 19 | NoHost, 20 | LimitExceeded(usize), 21 | Reqwest(reqwest::Error), 22 | } 23 | 24 | impl fmt::Display for DownloadError { 25 | #[rustfmt::skip] 26 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 27 | match self { 28 | DownloadError::ProxyNetworkError => write!(f, "Failed to connect to proxy"), 29 | DownloadError::InvalidStatus => write!(f, "Invalid status received from proxy"), 30 | DownloadError::LimitExceeded(limit) => write!(f, "The output file exceeded the maximum file size limit of {}. Try using a smaller input.", human_bytes((*limit) as f64)), 31 | DownloadError::Url(e) => write!(f, "Failed to parse URL: {e}"), 32 | DownloadError::NoHost => write!(f, "No host found in URL"), 33 | DownloadError::Reqwest(e) => write!(f, "{e}"), 34 | } 35 | } 36 | } 37 | 38 | impl std::error::Error for DownloadError {} 39 | 40 | fn get_next_proxy() -> &'static str { 41 | let config = &CONFIG; 42 | let len = config.urls.proxy.len(); 43 | 44 | (&config.urls.proxy[PROXY_NUM.fetch_add(1, Ordering::Relaxed) % len]) as _ 45 | } 46 | 47 | async fn download_with_proxy( 48 | client: &Client, 49 | url: &str, 50 | limit: usize, 51 | ) -> Result>, DownloadError> { 52 | let resp = client 53 | .get(format!("{}/proxy", get_next_proxy())) 54 | .header("User-Agent", "Assyst Discord Bot (https://github.com/jacherr/assyst2)") 55 | .query(&[("url", url), ("limit", &limit.to_string())]) 56 | .timeout(Duration::from_secs(10)) 57 | .send() 58 | .await 59 | .map_err(|_| DownloadError::ProxyNetworkError)?; 60 | 61 | if resp.status() != StatusCode::OK { 62 | return Err(DownloadError::InvalidStatus); 63 | } 64 | 65 | Ok(resp.bytes_stream()) 66 | } 67 | 68 | async fn download_no_proxy( 69 | client: &Client, 70 | url: &str, 71 | ) -> Result>, DownloadError> { 72 | Ok(client 73 | .get(url) 74 | .header("User-Agent", "Assyst Discord Bot (https://github.com/jacherr/assyst2)") 75 | .send() 76 | .await 77 | .map_err(DownloadError::Reqwest)? 78 | .bytes_stream()) 79 | } 80 | 81 | async fn read_stream(mut stream: S, limit: usize) -> Result, DownloadError> 82 | where 83 | S: Stream> + Unpin, 84 | { 85 | let mut bytes = Vec::new(); 86 | 87 | while let Some(Ok(chunk)) = stream.next().await { 88 | if bytes.len() > limit { 89 | return Err(DownloadError::LimitExceeded(limit)); 90 | } 91 | 92 | bytes.extend(chunk); 93 | } 94 | 95 | Ok(bytes) 96 | } 97 | 98 | /// Attempts to download a resource from a URL. 99 | pub async fn download_content( 100 | client: &Client, 101 | url: &str, 102 | limit: usize, 103 | untrusted: bool, 104 | ) -> Result, DownloadError> { 105 | const WHITELISTED_DOMAINS: &[&str] = &[ 106 | "tenor.com", 107 | "jacher.io", 108 | "discordapp.com", 109 | "discordapp.net", 110 | "wuk.sh", 111 | "gyazo.com", 112 | "cdn.discordapp.com", 113 | "media.discordapp.net", 114 | "notsobot.com", 115 | "twimg.com", 116 | "cdninstagram.com", 117 | "imput.net", 118 | ]; 119 | 120 | let config = &CONFIG; 121 | 122 | let url_p = Url::parse(url).map_err(DownloadError::Url)?; 123 | let host = url_p.host_str().ok_or(DownloadError::NoHost)?; 124 | 125 | let is_whitelisted = WHITELISTED_DOMAINS.iter().any(|d| host.contains(d)); 126 | 127 | if !config.urls.proxy.is_empty() && !is_whitelisted && untrusted { 128 | // First, try to download with proxy 129 | let stream = download_with_proxy(client, url, limit).await; 130 | 131 | if let Ok(stream) = stream { 132 | return read_stream(stream, limit).await; 133 | } 134 | } 135 | 136 | // Conditions for downloading with no proxy: 137 | // - Proxy not configured, 138 | // - Proxy failed, 139 | // - Domain is whitelisted 140 | let stream = download_no_proxy(client, url).await?; 141 | read_stream(stream, limit).await 142 | } 143 | -------------------------------------------------------------------------------- /assyst-core/src/gateway_handler/event_handlers/channel_update.rs: -------------------------------------------------------------------------------- 1 | use tracing::debug; 2 | use twilight_model::gateway::payload::incoming::ChannelUpdate; 3 | 4 | use crate::assyst::ThreadSafeAssyst; 5 | 6 | pub fn handle(assyst: ThreadSafeAssyst, event: ChannelUpdate) { 7 | if let Some(nsfw) = event.nsfw { 8 | assyst 9 | .rest_cache_handler 10 | .update_channel_age_restricted_status(event.id.get(), nsfw); 11 | 12 | debug!("Updated channel {} age restricted status to {nsfw}", event.id.get()); 13 | }; 14 | } 15 | -------------------------------------------------------------------------------- /assyst-core/src/gateway_handler/event_handlers/entitlement_create.rs: -------------------------------------------------------------------------------- 1 | use assyst_common::err; 2 | use assyst_common::macros::handle_log; 3 | use assyst_database::model::active_guild_premium_entitlement::ActiveGuildPremiumEntitlement; 4 | use tracing::info; 5 | use twilight_model::gateway::payload::incoming::EntitlementCreate; 6 | use twilight_model::guild::Guild; 7 | use twilight_model::id::marker::GuildMarker; 8 | use twilight_model::id::Id; 9 | 10 | use crate::assyst::ThreadSafeAssyst; 11 | 12 | pub async fn handle(assyst: ThreadSafeAssyst, event: EntitlementCreate) { 13 | let existing = assyst 14 | .entitlements 15 | .lock() 16 | .unwrap() 17 | .contains_key(&(event.id.get() as i64)); 18 | 19 | if existing { 20 | err!( 21 | "Entitlement ID {} (guild {:?} user {:?}) was created but already exists!", 22 | event.id, 23 | event.guild_id, 24 | event.user_id 25 | ); 26 | 27 | return; 28 | } 29 | 30 | // no expiry/created = test entitlement, requires special handling 31 | let active = match ActiveGuildPremiumEntitlement::try_from(event.0) { 32 | Err(e) => { 33 | err!("Error handling new entitlement: {e:?}"); 34 | return; 35 | }, 36 | Ok(a) => a, 37 | }; 38 | 39 | match active.set(&assyst.database_handler).await { 40 | Err(e) => { 41 | err!("Error registering new entitlement {}: {e:?}", active.entitlement_id); 42 | }, 43 | _ => {}, 44 | }; 45 | 46 | let expiry = active.expiry_unix_ms; 47 | let guild_id = Id::::new(active.guild_id as u64); 48 | let entitlement_id = active.entitlement_id; 49 | 50 | assyst.entitlements.lock().unwrap().insert(active.guild_id, active); 51 | 52 | let g: anyhow::Result = match assyst.http_client.guild(guild_id).await { 53 | Ok(g) => g.model().await.map_err(std::convert::Into::into), 54 | Err(e) => Err(e.into()), 55 | }; 56 | 57 | match g { 58 | Ok(g) => { 59 | handle_log(format!("New entitlement! Guild: {guild_id} - {}", g.name)); 60 | }, 61 | Err(_) => { 62 | handle_log(format!("New entitlement! Guild: {guild_id}")); 63 | }, 64 | } 65 | 66 | info!("Registered new entitlement: {entitlement_id} for guild {guild_id} (expiry unix {expiry})",); 67 | } 68 | -------------------------------------------------------------------------------- /assyst-core/src/gateway_handler/event_handlers/entitlement_delete.rs: -------------------------------------------------------------------------------- 1 | use assyst_common::err; 2 | use assyst_database::model::active_guild_premium_entitlement::ActiveGuildPremiumEntitlement; 3 | use twilight_model::gateway::payload::incoming::EntitlementDelete; 4 | 5 | use crate::assyst::ThreadSafeAssyst; 6 | 7 | pub async fn handle(assyst: ThreadSafeAssyst, event: EntitlementDelete) { 8 | let Some(guild_id) = event.guild_id else { 9 | err!( 10 | "Deleted entitlement ID {} (guild {:?} user {:?}) has no associated guild!", 11 | event.id, 12 | event.guild_id, 13 | event.user_id 14 | ); 15 | 16 | return; 17 | }; 18 | 19 | if event.user_id.is_none() { 20 | err!( 21 | "Deleted entitlement ID {} (guild {:?} user {:?}) has no associated user!", 22 | event.id, 23 | event.guild_id, 24 | event.user_id 25 | ); 26 | 27 | return; 28 | }; 29 | 30 | assyst.entitlements.lock().unwrap().remove(&(guild_id.get() as i64)); 31 | match ActiveGuildPremiumEntitlement::delete(&assyst.database_handler, event.id.get() as i64).await { 32 | Err(e) => { 33 | err!("Error deleting existing entitlement {}: {e:?}", event.id); 34 | }, 35 | _ => {}, 36 | } 37 | } 38 | -------------------------------------------------------------------------------- /assyst-core/src/gateway_handler/event_handlers/entitlement_update.rs: -------------------------------------------------------------------------------- 1 | use assyst_common::err; 2 | use assyst_common::macros::handle_log; 3 | use assyst_database::model::active_guild_premium_entitlement::ActiveGuildPremiumEntitlement; 4 | use tracing::info; 5 | use twilight_model::gateway::payload::incoming::EntitlementUpdate; 6 | use twilight_model::guild::Guild; 7 | use twilight_model::id::marker::GuildMarker; 8 | use twilight_model::id::Id; 9 | 10 | use crate::assyst::ThreadSafeAssyst; 11 | 12 | pub async fn handle(assyst: ThreadSafeAssyst, event: EntitlementUpdate) { 13 | // no expiry/created = test entitlement, requires special handling 14 | let new = match ActiveGuildPremiumEntitlement::try_from(event.0) { 15 | Err(e) => { 16 | err!("Error handling new entitlement: {e:?}"); 17 | return; 18 | }, 19 | Ok(a) => a, 20 | }; 21 | let guild_id = Id::::new(new.guild_id as u64); 22 | let entitlement_id = new.entitlement_id; 23 | 24 | match new.update(&assyst.database_handler).await { 25 | Err(e) => { 26 | err!("Error updating existing entitlement {entitlement_id}: {e:?}"); 27 | }, 28 | _ => {}, 29 | }; 30 | 31 | let g: anyhow::Result = match assyst.http_client.guild(guild_id).await { 32 | Ok(g) => g.model().await.map_err(std::convert::Into::into), 33 | Err(e) => Err(e.into()), 34 | }; 35 | 36 | match g { 37 | Ok(g) => { 38 | handle_log(format!("Updated entitlement! Guild: {guild_id} - {}", g.name)); 39 | }, 40 | Err(_) => { 41 | handle_log(format!("Updated entitlement! Guild: {guild_id}")); 42 | }, 43 | } 44 | 45 | assyst.entitlements.lock().unwrap().insert(guild_id.get() as i64, new); 46 | 47 | info!("Updated entitlement: {entitlement_id} for guild {guild_id}"); 48 | } 49 | -------------------------------------------------------------------------------- /assyst-core/src/gateway_handler/event_handlers/guild_create.rs: -------------------------------------------------------------------------------- 1 | use assyst_common::err; 2 | use tracing::info; 3 | use twilight_model::gateway::payload::incoming::GuildCreate; 4 | 5 | use crate::assyst::ThreadSafeAssyst; 6 | 7 | pub async fn handle(assyst: ThreadSafeAssyst, event: GuildCreate) { 8 | let id = event.id().get(); 9 | let (name, member_count) = match &event { 10 | GuildCreate::Available(g) => (g.name.clone(), g.member_count.unwrap_or(0)), 11 | GuildCreate::Unavailable(_) => (String::new(), 0), 12 | }; 13 | 14 | let should_handle = match assyst.persistent_cache_handler.handle_guild_create_event(event).await { 15 | Ok(s) => s, 16 | Err(e) => { 17 | err!("assyst-cache failed to handle GUILD_CREATE event: {}", e); 18 | return; 19 | }, 20 | }; 21 | 22 | if should_handle { 23 | info!("Joined guild {}: {} ({} members)", id, name, member_count); 24 | assyst.metrics_handler.inc_guilds(); 25 | } 26 | } 27 | -------------------------------------------------------------------------------- /assyst-core/src/gateway_handler/event_handlers/guild_delete.rs: -------------------------------------------------------------------------------- 1 | use assyst_common::err; 2 | use tracing::info; 3 | use twilight_model::gateway::payload::incoming::GuildDelete; 4 | 5 | use crate::assyst::ThreadSafeAssyst; 6 | 7 | pub async fn handle(assyst: ThreadSafeAssyst, event: GuildDelete) { 8 | let id = event.id.get(); 9 | 10 | let should_handle = match assyst.persistent_cache_handler.handle_guild_delete_event(event).await { 11 | Ok(s) => s, 12 | Err(e) => { 13 | err!("assyst-cache failed to handle GUILD_DELETE event: {}", e.to_string()); 14 | return; 15 | }, 16 | }; 17 | 18 | if should_handle { 19 | info!("Removed from guild {}", id); 20 | assyst.metrics_handler.dec_guilds(); 21 | } 22 | } 23 | -------------------------------------------------------------------------------- /assyst-core/src/gateway_handler/event_handlers/guild_update.rs: -------------------------------------------------------------------------------- 1 | use tracing::debug; 2 | use twilight_model::gateway::payload::incoming::GuildUpdate; 3 | 4 | use crate::assyst::ThreadSafeAssyst; 5 | 6 | pub fn handle(assyst: ThreadSafeAssyst, event: GuildUpdate) { 7 | assyst 8 | .rest_cache_handler 9 | .set_guild_owner(event.id.get(), event.owner_id.get()); 10 | 11 | assyst 12 | .rest_cache_handler 13 | .set_guild_upload_limit_bytes(event.id.get(), event.premium_tier); 14 | 15 | debug!("Updated guild {} cache info", event.id.get()); 16 | } 17 | -------------------------------------------------------------------------------- /assyst-core/src/gateway_handler/event_handlers/message_create.rs: -------------------------------------------------------------------------------- 1 | use std::collections::HashMap; 2 | use std::time::Instant; 3 | 4 | use assyst_common::err; 5 | use tracing::{debug, error}; 6 | use twilight_model::gateway::payload::incoming::MessageCreate; 7 | 8 | use super::after_command_execution_success; 9 | use crate::command::errors::{ExecutionError, TagParseError}; 10 | use crate::command::source::Source; 11 | use crate::command::{CommandCtxt, CommandData, RawMessageParseCtxt}; 12 | use crate::gateway_handler::message_parser::error::{ErrorSeverity, GetErrorSeverity}; 13 | use crate::gateway_handler::message_parser::parser::parse_message_into_command; 14 | use crate::ThreadSafeAssyst; 15 | 16 | /// Handle a [`MessageCreate`] event received from the Discord gateway. 17 | /// 18 | /// This function passes the message to the command parser, which then attempts to convert the 19 | /// message to a command for further processing. 20 | pub async fn handle(assyst: ThreadSafeAssyst, MessageCreate(message): MessageCreate) { 21 | if assyst.bad_translator.is_channel(message.channel_id.get()).await && !assyst.bad_translator.is_disabled().await { 22 | match assyst.bad_translator.handle_message(&assyst, Box::new(message)).await { 23 | Err(e) => { 24 | error!("BadTranslator channel execution failed: {e:?}"); 25 | }, 26 | _ => {}, 27 | }; 28 | return; 29 | } 30 | 31 | let processing_time_start = Instant::now(); 32 | 33 | match parse_message_into_command(assyst.clone(), &message, processing_time_start, false).await { 34 | Ok(Some(result)) => { 35 | let data = CommandData { 36 | source: Source::RawMessage, 37 | assyst: &assyst, 38 | execution_timings: result.execution_timings, 39 | calling_prefix: result.calling_prefix, 40 | message: Some(&message), 41 | interaction_subcommand: None, 42 | channel_id: message.channel_id, 43 | guild_id: message.guild_id, 44 | author: message.author.clone(), 45 | interaction_token: None, 46 | interaction_id: None, 47 | interaction_attachments: HashMap::new(), 48 | command_from_install_context: false, 49 | resolved_messages: None, 50 | resolved_users: None, 51 | }; 52 | let ctxt = RawMessageParseCtxt::new(CommandCtxt::new(&data), result.args); 53 | 54 | if let Err(err) = result.command.execute_raw_message(ctxt.clone()).await { 55 | match err.get_severity() { 56 | ErrorSeverity::Low => debug!("{err:?}"), 57 | ErrorSeverity::High => match err { 58 | // if invalid args: report usage to user 59 | ExecutionError::Parse(TagParseError::ArgsExhausted(_)) => { 60 | let _ = ctxt 61 | .cx 62 | .reply(format!( 63 | ":warning: `{err}\nUsage: {}{} {}`", 64 | ctxt.cx.data.calling_prefix, 65 | result.command.metadata().name, 66 | result.command.metadata().usage 67 | )) 68 | .await; 69 | }, 70 | _ => { 71 | let _ = ctxt.cx.reply(format!(":warning: ``{err:#}``")).await; 72 | }, 73 | }, 74 | } 75 | } else { 76 | let _ = after_command_execution_success(ctxt.cx, result.command) 77 | .await 78 | .map_err(|e| err!("Error handling post-command: {e:#}")); 79 | } 80 | }, 81 | Ok(None) => { /* command not found */ }, 82 | Err(error) => { 83 | if error.get_severity() == ErrorSeverity::High { 84 | err!("{error}"); 85 | } else { 86 | debug!("{error}"); 87 | } 88 | }, 89 | }; 90 | } 91 | -------------------------------------------------------------------------------- /assyst-core/src/gateway_handler/event_handlers/message_delete.rs: -------------------------------------------------------------------------------- 1 | use twilight_model::gateway::payload::incoming::MessageDelete; 2 | use twilight_model::id::Id; 3 | 4 | use crate::assyst::ThreadSafeAssyst; 5 | use crate::replies::ReplyState; 6 | 7 | /// Handle a [`MessageDelete`] event received from the Discord gateway. 8 | /// 9 | /// This function checks if the deleted message was one that invoked an Assyst command. 10 | /// If it was, then Assyst will attempt to delete the response to that command, to prevent any 11 | /// "dangling responses". 12 | pub async fn handle(assyst: ThreadSafeAssyst, message: MessageDelete) { 13 | if let Some(reply) = assyst.replies.get_raw_message(message.id.get()) 14 | && let ReplyState::InUse(reply) = reply.state 15 | { 16 | // ignore error 17 | _ = assyst 18 | .http_client 19 | .delete_message(message.channel_id, Id::new(reply.message_id)) 20 | .await; 21 | 22 | assyst.replies.remove_raw_message(message.id.get()); 23 | } 24 | } 25 | -------------------------------------------------------------------------------- /assyst-core/src/gateway_handler/event_handlers/mod.rs: -------------------------------------------------------------------------------- 1 | use assyst_database::model::command_usage::CommandUsage; 2 | 3 | use crate::command::{CommandCtxt, TCommand}; 4 | 5 | pub mod channel_update; 6 | pub mod entitlement_create; 7 | pub mod entitlement_delete; 8 | pub mod entitlement_update; 9 | pub mod guild_create; 10 | pub mod guild_delete; 11 | pub mod guild_update; 12 | pub mod interaction_create; 13 | pub mod message_create; 14 | pub mod message_delete; 15 | pub mod message_update; 16 | pub mod ready; 17 | 18 | pub async fn after_command_execution_success(ctxt: CommandCtxt<'_>, command: TCommand) -> anyhow::Result<()> { 19 | ctxt.assyst().metrics_handler.add_command(); 20 | ctxt.assyst() 21 | .metrics_handler 22 | .add_individual_command_usage(command.metadata().name) 23 | .await; 24 | (CommandUsage { 25 | command_name: command.metadata().name.to_owned(), 26 | uses: 0, 27 | }) 28 | .increment_command_uses(&ctxt.assyst().database_handler) 29 | .await?; 30 | 31 | Ok(()) 32 | } 33 | -------------------------------------------------------------------------------- /assyst-core/src/gateway_handler/event_handlers/ready.rs: -------------------------------------------------------------------------------- 1 | use assyst_common::config::CONFIG; 2 | use assyst_common::err; 3 | use assyst_string_fmt::Ansi; 4 | use tracing::{error, info}; 5 | use twilight_model::gateway::payload::incoming::Ready; 6 | use twilight_model::id::marker::ChannelMarker; 7 | use twilight_model::id::Id; 8 | 9 | use crate::assyst::ThreadSafeAssyst; 10 | 11 | /// Handle a shard sending a READY event. 12 | /// 13 | /// READY events are not particularly interesting, but it can be useful to see if any shards are 14 | /// resetting often. In addition, it provides a good gauge as to how much of the bot has started up, 15 | /// after a gateway restart. 16 | pub async fn handle(assyst: ThreadSafeAssyst, event: Ready) { 17 | if let Some(shard) = event.shard { 18 | info!( 19 | "Shard {} (total {}): {} in {} guilds", 20 | shard.number(), 21 | shard.total(), 22 | "READY".fg_green(), 23 | event.guilds.len() 24 | ); 25 | } 26 | 27 | if event.guilds.iter().any(|x| x.id.get() == CONFIG.dev.dev_guild) && CONFIG.dev.dev_message { 28 | let channel = Id::::new(CONFIG.dev.dev_channel); 29 | let _ = assyst 30 | .http_client 31 | .create_message(channel) 32 | .content("Dev shard is READY!") 33 | .await 34 | .inspect_err(|e| error!("FAILED to send shard ready message: {}", e.to_string())); 35 | } 36 | 37 | match assyst.persistent_cache_handler.handle_ready_event(event).await { 38 | Ok(num) => { 39 | info!("Adding {num} guilds to prometheus metrics from READY event"); 40 | assyst.metrics_handler.add_guilds(num); 41 | }, 42 | Err(e) => { 43 | err!("assyst-cache failed to handle READY event: {}", e.to_string()); 44 | }, 45 | } 46 | } 47 | -------------------------------------------------------------------------------- /assyst-core/src/gateway_handler/incoming_event.rs: -------------------------------------------------------------------------------- 1 | use twilight_model::gateway::event::{DispatchEvent, GatewayEvent}; 2 | use twilight_model::gateway::payload::incoming::{ 3 | ChannelUpdate, EntitlementCreate, EntitlementDelete, EntitlementUpdate, GuildCreate, GuildDelete, GuildUpdate, 4 | InteractionCreate, MessageCreate, MessageDelete, MessageUpdate, Ready, 5 | }; 6 | 7 | #[derive(Debug)] 8 | pub enum IncomingEvent { 9 | ChannelUpdate(ChannelUpdate), 10 | GuildCreate(Box), // this struct is huge. 11 | GuildDelete(GuildDelete), 12 | GuildUpdate(GuildUpdate), 13 | InteractionCreate(Box), 14 | MessageCreate(Box), // same problem 15 | MessageDelete(MessageDelete), 16 | MessageUpdate(MessageUpdate), 17 | ShardReady(Ready), 18 | EntitlementCreate(EntitlementCreate), 19 | EntitlementUpdate(EntitlementUpdate), 20 | EntitlementDelete(EntitlementDelete), 21 | } 22 | impl TryFrom for IncomingEvent { 23 | type Error = (); 24 | 25 | fn try_from(event: GatewayEvent) -> Result { 26 | match event { 27 | GatewayEvent::Dispatch(_, event) => match event { 28 | DispatchEvent::MessageCreate(message) => Ok(IncomingEvent::MessageCreate(message)), 29 | DispatchEvent::MessageUpdate(message) => Ok(IncomingEvent::MessageUpdate(*message)), 30 | DispatchEvent::MessageDelete(message) => Ok(IncomingEvent::MessageDelete(message)), 31 | DispatchEvent::GuildCreate(guild) => Ok(IncomingEvent::GuildCreate(guild)), 32 | DispatchEvent::GuildDelete(guild) => Ok(IncomingEvent::GuildDelete(guild)), 33 | DispatchEvent::GuildUpdate(guild) => Ok(IncomingEvent::GuildUpdate(*guild)), 34 | DispatchEvent::Ready(ready) => Ok(IncomingEvent::ShardReady(*ready)), 35 | DispatchEvent::ChannelUpdate(channel) => Ok(IncomingEvent::ChannelUpdate(*channel)), 36 | DispatchEvent::InteractionCreate(interaction) => Ok(IncomingEvent::InteractionCreate(interaction)), 37 | DispatchEvent::EntitlementCreate(e) => Ok(IncomingEvent::EntitlementCreate(e)), 38 | DispatchEvent::EntitlementUpdate(e) => Ok(IncomingEvent::EntitlementUpdate(e)), 39 | DispatchEvent::EntitlementDelete(e) => Ok(IncomingEvent::EntitlementDelete(e)), 40 | _ => Err(()), 41 | }, 42 | _ => Err(()), 43 | } 44 | } 45 | } 46 | -------------------------------------------------------------------------------- /assyst-core/src/gateway_handler/message_parser/error.rs: -------------------------------------------------------------------------------- 1 | use std::fmt::Display; 2 | 3 | use twilight_model::channel::message::MessageType; 4 | 5 | pub trait GetErrorSeverity { 6 | fn get_severity(&self) -> ErrorSeverity; 7 | } 8 | 9 | #[derive(Debug)] 10 | /// An error when pre-processing the message. 11 | pub enum PreParseError { 12 | /// Message does not start with the correct prefix. 13 | MessageNotPrefixed(String), 14 | /// Invocating user is globally blacklisted from using the bot. 15 | UserGloballyBlacklisted(u64), 16 | /// Invocating user is a bot or webhook. 17 | UserIsBotOrWebhook(Option), 18 | /// The kind of message is not supported, e.g., a user join system message. 19 | UnsupportedMessageKind(MessageType), 20 | /// A `MESSAGE_UPDATE` was received, but it had no edited timestamp. 21 | EditedMessageWithNoTimestamp, 22 | /// Other unknown failure. Unexpected error with high severity. 23 | Failure(String), 24 | } 25 | impl Display for PreParseError { 26 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 27 | match self { 28 | Self::MessageNotPrefixed(prefix) => { 29 | write!(f, "Message does not start with correct prefix ({prefix})") 30 | }, 31 | Self::UserGloballyBlacklisted(id) => { 32 | write!(f, "User {id} is globally blacklisted") 33 | }, 34 | Self::UserIsBotOrWebhook(id) => { 35 | write!(f, "User is a bot or webhook ({})", id.unwrap_or(0)) 36 | }, 37 | Self::UnsupportedMessageKind(kind) => { 38 | write!(f, "Unsupported message kind ({kind:?})") 39 | }, 40 | Self::EditedMessageWithNoTimestamp => f.write_str("The message was updated, but not edited."), 41 | Self::Failure(message) => { 42 | write!(f, "Preprocessor failure: {message}") 43 | }, 44 | } 45 | } 46 | } 47 | impl GetErrorSeverity for PreParseError { 48 | fn get_severity(&self) -> ErrorSeverity { 49 | match self { 50 | PreParseError::Failure(_) => ErrorSeverity::High, 51 | _ => ErrorSeverity::Low, 52 | } 53 | } 54 | } 55 | impl std::error::Error for PreParseError {} 56 | 57 | #[derive(Debug)] 58 | pub enum MetadataCheckInvalidated {} 59 | impl GetErrorSeverity for MetadataCheckInvalidated { 60 | fn get_severity(&self) -> ErrorSeverity { 61 | match self { 62 | _ => todo!(), 63 | } 64 | } 65 | } 66 | impl Display for MetadataCheckInvalidated { 67 | fn fmt(&self, _f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 68 | match self { 69 | _ => todo!(), 70 | } 71 | } 72 | } 73 | 74 | #[derive(Debug)] 75 | pub enum ParseError { 76 | /// Failure with preprocessing of the message. 77 | PreParseFail(PreParseError), 78 | } 79 | impl Display for ParseError { 80 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 81 | match self { 82 | Self::PreParseFail(message) => { 83 | write!(f, "Pre-parse failed: {message}") 84 | }, 85 | } 86 | } 87 | } 88 | impl std::error::Error for ParseError {} 89 | impl GetErrorSeverity for ParseError { 90 | fn get_severity(&self) -> ErrorSeverity { 91 | match self { 92 | ParseError::PreParseFail(e) => e.get_severity(), 93 | } 94 | } 95 | } 96 | impl From for ParseError { 97 | fn from(value: PreParseError) -> Self { 98 | ParseError::PreParseFail(value) 99 | } 100 | } 101 | 102 | #[derive(PartialEq, Eq)] 103 | pub enum ErrorSeverity { 104 | Low, 105 | High, 106 | } 107 | -------------------------------------------------------------------------------- /assyst-core/src/gateway_handler/message_parser/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod error; 2 | pub mod parser; 3 | pub mod preprocess; 4 | -------------------------------------------------------------------------------- /assyst-core/src/gateway_handler/message_parser/parser.rs: -------------------------------------------------------------------------------- 1 | use std::time::Instant; 2 | 3 | use twilight_model::channel::Message; 4 | 5 | use super::error::ParseError; 6 | use super::preprocess::preprocess; 7 | use crate::command::registry::find_command_by_name; 8 | use crate::command::{ExecutionTimings, TCommand}; 9 | use crate::ThreadSafeAssyst; 10 | 11 | pub struct ParseResult<'a> { 12 | pub command: TCommand, 13 | pub args: &'a str, 14 | pub calling_prefix: String, 15 | pub execution_timings: ExecutionTimings, 16 | } 17 | 18 | /// Parse any generic Message object into a Command. 19 | /// 20 | /// This function takes all steps necessary to split a message into critical command components, 21 | /// and if at any point the parse fails, then return with no action. 22 | /// 23 | /// After parsing, a `CoreEvent` is fired to assyst-core signaling that the command should be 24 | /// executed. Parsing a message has several steps.
25 | /// **Step 1**: Check if the invocating user is blacklisted. If so, prematurely return. 26 | /// 27 | /// **Step 2**: Check that the message starts with the correct prefix. 28 | /// The prefix can be one of four things: 29 | /// 1. The guild-specific prefix, stored in the database, 30 | /// 2. No prefix, if the command is ran in DMs, 31 | /// 3. The bot's mention, in the form of @Assyst, 32 | /// 4. The prefix override, if specified, in config.toml. 33 | /// The mention prefix takes precedence over all other, followed by the prefix override, 34 | /// followed by the guild prefix. 35 | /// This function identifies the prefix and checks if it is valid for this particular invocation. 36 | /// If it is not, then prematurely return. 37 | /// 38 | /// **Step 3**: Check if this Message already has an associated reply (if, for example, the 39 | /// invocation was updated). 40 | /// These events have a timeout for handling, to prevent editing of very old 41 | /// messages. If it is expired, prematurely return. 42 | /// 43 | /// **Step 4**: Parse the Command from the Message itself. If it fails to parse, prematurely return. 44 | /// 45 | /// Once all steps are complete, a Command is returned, ready for execution. 46 | /// Note that metadata is checked *during* execution (i.e., in the base command's `Command::execute` 47 | /// implementation, see [`crate::command::check_metadata`]) 48 | pub async fn parse_message_into_command( 49 | assyst: ThreadSafeAssyst, 50 | message: &Message, 51 | processing_time_start: Instant, 52 | from_edit: bool, 53 | ) -> Result, ParseError> { 54 | let parse_start = Instant::now(); 55 | let preprocess_start = Instant::now(); 56 | 57 | let preprocess = preprocess(assyst.clone(), message, from_edit).await?; 58 | 59 | let preprocess_time = preprocess_start.elapsed(); 60 | 61 | // commands can theoretically have spaces in their name so we need to try and identify the 62 | // set of 'words' in source text to associate with a command name (essentially finding the 63 | // divide between command name and command arguments) 64 | let command_text = &message.content[preprocess.prefix.len()..]; 65 | 66 | let mut args = command_text.split_ascii_whitespace(); 67 | let Some(command) = args.next() else { 68 | return Ok(None); 69 | }; 70 | let args = args.remainder().unwrap_or(""); 71 | let Some(command) = find_command_by_name(command) else { 72 | return Ok(None); 73 | }; 74 | 75 | Ok(Some(ParseResult { 76 | command, 77 | args, 78 | calling_prefix: preprocess.prefix, 79 | execution_timings: ExecutionTimings { 80 | processing_time_start, 81 | parse_total: parse_start.elapsed(), 82 | prefix_determiner: preprocess.prefixing_determinism_time, 83 | metadata_check_start: Instant::now(), 84 | preprocess_total: preprocess_time, 85 | }, 86 | })) 87 | } 88 | -------------------------------------------------------------------------------- /assyst-core/src/gateway_handler/mod.rs: -------------------------------------------------------------------------------- 1 | #![allow(clippy::match_single_binding, clippy::single_match)] // shh... 2 | use self::incoming_event::IncomingEvent; 3 | use crate::assyst::ThreadSafeAssyst; 4 | 5 | pub mod event_handlers; 6 | pub mod incoming_event; 7 | pub mod message_parser; 8 | pub mod reply; 9 | 10 | /// Checks the enum variant of this `IncomingEvent` and calls the appropriate handler function 11 | /// for further processing. 12 | pub async fn handle_raw_event(context: ThreadSafeAssyst, event: IncomingEvent) { 13 | match event { 14 | IncomingEvent::ShardReady(event) => { 15 | event_handlers::ready::handle(context, event).await; 16 | }, 17 | IncomingEvent::MessageCreate(event) => { 18 | event_handlers::message_create::handle(context, *event).await; 19 | }, 20 | IncomingEvent::MessageUpdate(event) => { 21 | event_handlers::message_update::handle(context, event).await; 22 | }, 23 | IncomingEvent::MessageDelete(event) => { 24 | event_handlers::message_delete::handle(context, event).await; 25 | }, 26 | IncomingEvent::GuildCreate(event) => { 27 | event_handlers::guild_create::handle(context, *event).await; 28 | }, 29 | IncomingEvent::GuildDelete(event) => { 30 | event_handlers::guild_delete::handle(context, event).await; 31 | }, 32 | IncomingEvent::GuildUpdate(event) => { 33 | event_handlers::guild_update::handle(context, event); 34 | }, 35 | IncomingEvent::ChannelUpdate(event) => { 36 | event_handlers::channel_update::handle(context, event); 37 | }, 38 | IncomingEvent::InteractionCreate(event) => { 39 | event_handlers::interaction_create::handle(context, *event).await; 40 | }, 41 | IncomingEvent::EntitlementCreate(event) => { 42 | event_handlers::entitlement_create::handle(context, event).await; 43 | }, 44 | IncomingEvent::EntitlementUpdate(event) => { 45 | event_handlers::entitlement_update::handle(context, event).await; 46 | }, 47 | IncomingEvent::EntitlementDelete(event) => { 48 | event_handlers::entitlement_delete::handle(context, event).await; 49 | }, 50 | } 51 | } 52 | -------------------------------------------------------------------------------- /assyst-core/src/persistent_cache_handler/mod.rs: -------------------------------------------------------------------------------- 1 | use assyst_common::cache::{ 2 | CacheJob, CacheJobSend, CacheResponse, CacheResponseSend, GuildCreateData, GuildDeleteData, ReadyData, 3 | }; 4 | use assyst_common::pipe::Pipe; 5 | use assyst_common::unwrap_enum_variant; 6 | use tokio::spawn; 7 | use tokio::sync::mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender}; 8 | use tokio::sync::oneshot; 9 | use tracing::{info, warn}; 10 | use twilight_model::gateway::payload::incoming::{GuildCreate, GuildDelete, Ready}; 11 | 12 | /// Handles communication with assyst-cache. 13 | pub struct PersistentCacheHandler { 14 | pub cache_tx: UnboundedSender, 15 | } 16 | impl PersistentCacheHandler { 17 | pub fn new(path: &str) -> PersistentCacheHandler { 18 | let (tx, rx) = unbounded_channel::(); 19 | PersistentCacheHandler::init_pipe(rx, path); 20 | PersistentCacheHandler { cache_tx: tx } 21 | } 22 | 23 | fn init_pipe(mut rx: UnboundedReceiver, path: &str) { 24 | let path = path.to_owned(); 25 | // main handler thread 26 | spawn(async move { 27 | info!("Connecting to assyst-cache pipe on {path}"); 28 | loop { 29 | let mut pipe = Pipe::poll_connect(&path, None).await.unwrap(); 30 | info!("Connected to assyst-cache pipe on {path}"); 31 | loop { 32 | // ok to unwrap because tx is permanently stored in handler 33 | let (tx, data) = rx.recv().await.unwrap(); 34 | 35 | if let Err(e) = pipe.write_object(data).await { 36 | // safe to unwrap because no situation in which the channel should be 37 | // dropped 38 | tx.send(Err(e)).unwrap(); 39 | break; 40 | }; 41 | 42 | let result = match pipe.read_object::().await { 43 | Ok(x) => x, 44 | Err(e) => { 45 | tx.send(Err(e)).unwrap(); 46 | break; 47 | }, 48 | }; 49 | 50 | tx.send(Ok(result)).unwrap(); 51 | } 52 | warn!("Communication to assyst-cache lost, attempting reconnection"); 53 | } 54 | }); 55 | } 56 | 57 | async fn run_cache_job(&self, job: CacheJob) -> anyhow::Result { 58 | let (tx, rx) = oneshot::channel::(); 59 | // can unwrap since it should never close 60 | self.cache_tx.send((tx, job)).unwrap(); 61 | rx.await.unwrap() 62 | } 63 | 64 | /// Handles a READY event, caching its guilds. Returns the number of newly cached guilds. 65 | pub async fn handle_ready_event(&self, event: Ready) -> anyhow::Result { 66 | self.run_cache_job(CacheJob::HandleReady(ReadyData::from(event))) 67 | .await 68 | .map(|x| unwrap_enum_variant!(x, CacheResponse::NewGuildsFromReady)) 69 | } 70 | 71 | /// Handles a `GUILD_CREATE`. This method returns a bool which states if this guild is new or not. 72 | /// A new guild is one that was not received during the start-up of the gateway connection. 73 | pub async fn handle_guild_create_event(&self, event: GuildCreate) -> anyhow::Result { 74 | self.run_cache_job(CacheJob::HandleGuildCreate(GuildCreateData::from(event))) 75 | .await 76 | .map(|x| unwrap_enum_variant!(x, CacheResponse::ShouldHandleGuildCreate)) 77 | } 78 | 79 | /// Handles a `GUILD_DELETE`. This method returns a bool which states if the bot was actually 80 | /// kicked from this guild. 81 | pub async fn handle_guild_delete_event(&self, event: GuildDelete) -> anyhow::Result { 82 | self.run_cache_job(CacheJob::HandleGuildDelete(GuildDeleteData::from(event))) 83 | .await 84 | .map(|x| unwrap_enum_variant!(x, CacheResponse::ShouldHandleGuildDelete)) 85 | } 86 | 87 | pub async fn get_guild_count(&self) -> anyhow::Result { 88 | let request = CacheJob::GetGuildCount; 89 | let response = self.run_cache_job(request).await?; 90 | Ok(unwrap_enum_variant!(response, CacheResponse::TotalGuilds)) 91 | } 92 | } 93 | -------------------------------------------------------------------------------- /assyst-core/src/replies.rs: -------------------------------------------------------------------------------- 1 | use std::time::{Duration, Instant}; 2 | 3 | use moka::sync::Cache; 4 | 5 | #[derive(Copy, Clone, Debug)] 6 | pub struct ReplyInUse { 7 | /// The message ID of this reply 8 | pub message_id: u64, 9 | /// Whether the reply has any attachments. 10 | pub _has_attachments: bool, 11 | } 12 | 13 | #[allow(dead_code)] 14 | #[derive(Debug, Clone)] 15 | pub enum ReplyState { 16 | Processing, 17 | InUse(ReplyInUse), 18 | } 19 | 20 | #[derive(Debug, Clone)] 21 | pub struct Reply { 22 | pub state: ReplyState, 23 | pub _created: Instant, 24 | } 25 | 26 | impl Reply { 27 | pub fn in_use(&self) -> Option { 28 | if let ReplyState::InUse(reply) = self.state { 29 | Some(reply) 30 | } else { 31 | None 32 | } 33 | } 34 | } 35 | 36 | /// Cached command replies. First cache is for "raw" messages, second is for interaction messages. 37 | pub struct Replies(Cache, Cache); 38 | 39 | impl Replies { 40 | pub fn new() -> Self { 41 | Self( 42 | Cache::builder() 43 | .max_capacity(1000) 44 | .time_to_idle(Duration::from_secs(60 * 5)) 45 | .build(), 46 | Cache::builder() 47 | .max_capacity(1000) 48 | .time_to_idle(Duration::from_secs(60 * 5)) 49 | .build(), 50 | ) 51 | } 52 | 53 | pub fn insert_raw_message(&self, id: u64, reply: Reply) { 54 | self.0.insert(id, reply); 55 | } 56 | 57 | pub fn remove_raw_message(&self, id: u64) -> Option { 58 | self.0.remove(&id) 59 | } 60 | 61 | pub fn get_raw_message(&self, id: u64) -> Option { 62 | self.0.get(&id) 63 | } 64 | 65 | pub fn insert_interaction_command(&self, id: u64) { 66 | self.1.insert(id, ()); 67 | } 68 | 69 | pub fn get_interaction_command(&self, id: u64) -> Option<()> { 70 | self.1.get(&id) 71 | } 72 | } 73 | -------------------------------------------------------------------------------- /assyst-core/src/rest/audio_identification.rs: -------------------------------------------------------------------------------- 1 | use assyst_common::config::CONFIG; 2 | use reqwest::multipart::{Form, Part}; 3 | use reqwest::Client; 4 | 5 | const NOT_SO_IDENTIFY_URL: &str = "https://notsobot.com/api/media/av/tools/identify"; 6 | 7 | pub mod notsoidentify { 8 | use serde::{Deserialize, Serialize}; 9 | #[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)] 10 | pub struct Song { 11 | pub album: NamedField, 12 | pub artists: Vec, 13 | pub title: String, 14 | pub platforms: Platform, 15 | } 16 | 17 | #[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)] 18 | pub struct NamedField { 19 | pub name: String, 20 | } 21 | 22 | #[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)] 23 | pub struct Platform { 24 | pub youtube: Option, 25 | } 26 | 27 | #[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)] 28 | pub struct YouTube { 29 | pub url: String, 30 | } 31 | } 32 | 33 | pub async fn identify_song_notsoidentify(client: &Client, search: String) -> anyhow::Result> { 34 | let formdata = Form::new(); 35 | let formdata = formdata.part("url", Part::text(search)); 36 | Ok(client 37 | .post(NOT_SO_IDENTIFY_URL) 38 | .header("authorization", CONFIG.authentication.notsoapi.to_string()) 39 | .multipart(formdata) 40 | .send() 41 | .await? 42 | .error_for_status()? 43 | .json::>() 44 | .await?) 45 | } 46 | -------------------------------------------------------------------------------- /assyst-core/src/rest/bad_translation.rs: -------------------------------------------------------------------------------- 1 | use std::error::Error; 2 | use std::fmt::Display; 3 | 4 | use assyst_common::config::CONFIG; 5 | use reqwest::{Client, Error as ReqwestError}; 6 | use serde::Deserialize; 7 | 8 | const MAX_ATTEMPTS: u8 = 5; 9 | 10 | mod routes { 11 | pub const LANGUAGES: &str = "/languages"; 12 | } 13 | 14 | #[derive(Debug)] 15 | pub enum TranslateError { 16 | Reqwest(ReqwestError), 17 | Raw(&'static str), 18 | } 19 | 20 | impl Display for TranslateError { 21 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 22 | match self { 23 | TranslateError::Reqwest(e) => write!( 24 | f, 25 | "A network error occurred: {}", 26 | e.to_string().replace(&CONFIG.urls.bad_translation, "[bt]") 27 | ), 28 | TranslateError::Raw(s) => write!(f, "{s}"), 29 | } 30 | } 31 | } 32 | 33 | impl Error for TranslateError {} 34 | 35 | #[derive(Deserialize)] 36 | pub struct Translation { 37 | pub lang: String, 38 | pub text: String, 39 | } 40 | 41 | #[derive(Deserialize)] 42 | pub struct TranslateResult { 43 | pub translations: Vec, 44 | pub result: Translation, 45 | } 46 | 47 | async fn translate_retry( 48 | client: &Client, 49 | text: &str, 50 | target: Option<&str>, 51 | count: Option, 52 | additional_data: Option<&[(&str, String)]>, 53 | ) -> Result { 54 | let mut query_args = vec![("text", text.to_owned())]; 55 | 56 | if let Some(target) = target { 57 | query_args.push(("target", target.to_owned())); 58 | } 59 | 60 | if let Some(count) = count { 61 | query_args.push(("count", count.to_string())); 62 | } 63 | 64 | if let Some(data) = additional_data { 65 | for (k, v) in data { 66 | query_args.push((k, v.to_string())); 67 | } 68 | } 69 | 70 | client 71 | .get(&CONFIG.urls.bad_translation) 72 | .query(&query_args) 73 | .send() 74 | .await 75 | .map_err(TranslateError::Reqwest)? 76 | .json() 77 | .await 78 | .map_err(TranslateError::Reqwest) 79 | } 80 | 81 | async fn translate( 82 | client: &Client, 83 | text: &str, 84 | target: Option<&str>, 85 | count: Option, 86 | additional_data: Option<&[(&str, String)]>, 87 | ) -> Result { 88 | let mut attempt = 0; 89 | 90 | while attempt <= MAX_ATTEMPTS { 91 | match translate_retry(client, text, target, count, additional_data).await { 92 | Ok(result) => return Ok(result), 93 | Err(e) => eprintln!("Proxy failed! {e:?}"), 94 | }; 95 | 96 | attempt += 1; 97 | } 98 | 99 | Err(TranslateError::Raw("BT Failed: Too many attempts")) 100 | } 101 | 102 | pub async fn bad_translate(client: &Client, text: &str) -> Result { 103 | translate(client, text, None, None, None).await 104 | } 105 | 106 | pub async fn bad_translate_target( 107 | client: &Client, 108 | text: &str, 109 | target: &str, 110 | ) -> Result { 111 | translate(client, text, Some(target), None, None).await 112 | } 113 | 114 | pub async fn bad_translate_with_count( 115 | client: &Client, 116 | text: &str, 117 | count: u32, 118 | ) -> Result { 119 | translate(client, text, None, Some(count), None).await 120 | } 121 | 122 | pub async fn translate_single(client: &Client, text: &str, target: &str) -> Result { 123 | translate(client, text, Some(target), Some(1), None).await 124 | } 125 | 126 | pub async fn get_languages(client: &Client) -> Result, Box)>, TranslateError> { 127 | client 128 | .get(format!("{}{}", CONFIG.urls.bad_translation, routes::LANGUAGES)) 129 | .send() 130 | .await 131 | .map_err(TranslateError::Reqwest)? 132 | .json() 133 | .await 134 | .map_err(TranslateError::Reqwest) 135 | } 136 | 137 | // used in btchannel later 138 | #[allow(unused)] 139 | pub async fn validate_language(client: &Client, provided_language: &str) -> Result { 140 | let languages = get_languages(client).await?; 141 | Ok(languages.iter().any(|(language, _)| &**language == provided_language)) 142 | } 143 | -------------------------------------------------------------------------------- /assyst-core/src/rest/charinfo.rs: -------------------------------------------------------------------------------- 1 | use reqwest::Client; 2 | 3 | static CHARINFO_URL: &str = "https://www.fileformat.info/info/unicode/char/"; 4 | 5 | pub async fn get_char_info(client: &Client, ch: char) -> anyhow::Result<(String, String)> { 6 | let url = format!("{}{:x}", CHARINFO_URL, ch as u32); 7 | 8 | Ok((client.get(&url).send().await?.text().await?, url)) 9 | } 10 | 11 | /// Attempts to extract the page title for charingo 12 | pub fn extract_page_title(input: &str) -> Option { 13 | let dom = tl::parse(input, tl::ParserOptions::default()).ok()?; 14 | let parser = dom.parser(); 15 | 16 | let tag = dom.query_selector("title")?.next()?.get(parser)?; 17 | 18 | Some(tag.inner_text(parser).into_owned()) 19 | } 20 | -------------------------------------------------------------------------------- /assyst-core/src/rest/eval.rs: -------------------------------------------------------------------------------- 1 | use anyhow::Context; 2 | use assyst_common::config::CONFIG; 3 | use assyst_common::eval::{FakeEvalBody, FakeEvalImageResponse, FakeEvalMessageData}; 4 | use assyst_common::util::filetype::get_sig; 5 | use reqwest::Client; 6 | use twilight_model::channel::Message; 7 | 8 | pub async fn fake_eval( 9 | client: &Client, 10 | code: String, 11 | accept_image: bool, 12 | message: Option<&Message>, 13 | args: Vec, 14 | ) -> anyhow::Result { 15 | let response = client 16 | .post(format!("{}/eval", CONFIG.urls.eval)) 17 | .query(&[("returnBuffer", accept_image)]) 18 | .json(&FakeEvalBody { 19 | code, 20 | data: Some(FakeEvalMessageData { args, message }), 21 | }) 22 | .send() 23 | .await? 24 | .bytes() 25 | .await?; 26 | 27 | if let Some(sig) = get_sig(&response) { 28 | Ok(FakeEvalImageResponse::Image(response.to_vec(), sig)) 29 | } else { 30 | let text = std::str::from_utf8(&response).context("eval returned non-utf8 text response")?; 31 | Ok(FakeEvalImageResponse::Text(serde_json::from_str(text)?)) 32 | } 33 | } 34 | -------------------------------------------------------------------------------- /assyst-core/src/rest/filer.rs: -------------------------------------------------------------------------------- 1 | use anyhow::Context; 2 | use assyst_common::config::CONFIG; 3 | use reqwest::Client; 4 | use serde::Deserialize; 5 | 6 | #[derive(Deserialize)] 7 | pub struct FilerStats { 8 | pub count: u64, 9 | pub size_bytes: u64, 10 | } 11 | 12 | pub async fn get_filer_stats(client: &Client) -> anyhow::Result { 13 | Ok(client 14 | .get(format!("{}/stats", CONFIG.urls.filer)) 15 | .send() 16 | .await? 17 | .json::() 18 | .await?) 19 | } 20 | 21 | pub async fn upload_to_filer(client: &Client, data: Vec, content_type: &str) -> anyhow::Result { 22 | Ok(client 23 | .post(&CONFIG.urls.filer) 24 | .header(reqwest::header::CONTENT_TYPE, content_type) 25 | .header(reqwest::header::AUTHORIZATION, &CONFIG.authentication.filer_key) 26 | .body(data) 27 | .send() 28 | .await? 29 | .error_for_status() 30 | .context("Failed to upload to filer")? 31 | .text() 32 | .await?) 33 | } 34 | -------------------------------------------------------------------------------- /assyst-core/src/rest/identify.rs: -------------------------------------------------------------------------------- 1 | use assyst_common::config::CONFIG; 2 | use reqwest::Client; 3 | use serde::{Deserialize, Serialize}; 4 | 5 | const IDENTIFY_ROUTE: &str = "https://microsoft-computer-vision3.p.rapidapi.com/analyze?language=en&descriptionExclude=Celebrities&visualFeatures=Description&details=Celebrities"; 6 | 7 | #[derive(Serialize)] 8 | pub struct IdentifyBody<'a> { 9 | pub url: &'a str, 10 | } 11 | 12 | #[derive(Deserialize)] 13 | pub struct IdentifyResponse { 14 | pub description: Option, 15 | } 16 | 17 | #[derive(Deserialize)] 18 | pub struct IdentifyDescription { 19 | pub captions: Vec, 20 | } 21 | 22 | #[derive(Deserialize)] 23 | pub struct IdentifyCaption { 24 | pub text: String, 25 | pub confidence: f32, 26 | } 27 | 28 | pub async fn identify_image(client: &Client, url: &str) -> reqwest::Result { 29 | client 30 | .post(IDENTIFY_ROUTE) 31 | .header("x-rapidapi-key", &CONFIG.authentication.rapidapi_token) 32 | .json(&IdentifyBody { url }) 33 | .send() 34 | .await? 35 | .json() 36 | .await 37 | } 38 | -------------------------------------------------------------------------------- /assyst-core/src/rest/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod audio_identification; 2 | pub mod bad_translation; 3 | pub mod charinfo; 4 | pub mod cooltext; 5 | pub mod eval; 6 | pub mod filer; 7 | pub mod identify; 8 | pub mod patreon; 9 | pub mod r34; 10 | pub mod rest_cache_handler; 11 | pub mod rust; 12 | pub mod top_gg; 13 | pub mod web_media_download; 14 | 15 | pub static NORMAL_DISCORD_UPLOAD_LIMIT_BYTES: u64 = 10_000_000; 16 | pub static PREMIUM_TIER2_DISCORD_UPLOAD_LIMIT_BYTES: u64 = 50_000_000; 17 | pub static PREMIUM_TIER3_DISCORD_UPLOAD_LIMIT_BYTES: u64 = 100_000_000; 18 | -------------------------------------------------------------------------------- /assyst-core/src/rest/r34.rs: -------------------------------------------------------------------------------- 1 | use anyhow::{bail, Context}; 2 | use rand::prelude::SliceRandom; 3 | use reqwest::Client; 4 | use serde::Deserialize; 5 | 6 | static R34_URL: &str = "https://api.rule34.xxx/index.php?tags="; 7 | 8 | #[derive(Deserialize, Clone)] 9 | pub struct R34Result { 10 | pub file_url: String, 11 | pub score: i32, 12 | } 13 | 14 | pub async fn get_random_r34(client: &Client, tags: &str) -> anyhow::Result { 15 | let all = client 16 | .get(format!("{}{}", R34_URL, &tags.replace(' ', "+")[..])) 17 | .query(&[ 18 | ("page", "dapi"), 19 | ("s", "post"), 20 | ("q", "index"), 21 | ("json", "1"), 22 | ("limit", "1000"), 23 | ]) 24 | .send() 25 | .await? 26 | .error_for_status()? 27 | .json::>() 28 | .await; 29 | 30 | if let Err(e) = all { 31 | if e.to_string().contains("EOF") { 32 | bail!("No results found") 33 | } 34 | Err(e.into()) 35 | } else { 36 | let all = all.unwrap(); 37 | 38 | let mut rng = rand::thread_rng(); 39 | 40 | all.choose(&mut rng).cloned().context("No results found") 41 | } 42 | } 43 | -------------------------------------------------------------------------------- /assyst-core/src/rest/top_gg.rs: -------------------------------------------------------------------------------- 1 | use std::sync::LazyLock; 2 | 3 | use assyst_common::config::CONFIG; 4 | use reqwest::Client; 5 | use serde_json::json; 6 | 7 | static ROUTE: LazyLock = LazyLock::new(|| format!("https://top.gg/api/bots/{}/stats", CONFIG.bot_id)); 8 | 9 | pub async fn post_top_gg_stats(client: &Client, guild_count: u64) -> anyhow::Result<()> { 10 | client 11 | .post(&*ROUTE) 12 | .header("authorization", &CONFIG.authentication.top_gg_token) 13 | .json(&json!({ "server_count": guild_count })) 14 | .send() 15 | .await? 16 | .error_for_status()?; 17 | 18 | Ok(()) 19 | } 20 | -------------------------------------------------------------------------------- /assyst-core/src/task/mod.rs: -------------------------------------------------------------------------------- 1 | use std::future::Future; 2 | use std::pin::Pin; 3 | use std::time::Duration; 4 | 5 | use tokio::spawn; 6 | use tokio::task::JoinHandle; 7 | use tokio::time::sleep; 8 | 9 | use crate::assyst::ThreadSafeAssyst; 10 | 11 | pub type TaskResult = Pin + Send>>; 12 | pub type TaskRun = Box TaskResult + Send + Sync>; 13 | 14 | pub mod tasks; 15 | 16 | #[macro_export] 17 | macro_rules! function_task_callback { 18 | ($expression:expr) => { 19 | Box::new(move |assyst: ThreadSafeAssyst| Box::pin($expression(assyst.clone()))) 20 | }; 21 | } 22 | 23 | /// A Task is a function which is called repeatedly on a set interval. 24 | /// 25 | /// A Task can be created to run on its own thread, and once per interval the provided function will 26 | /// be executed. 27 | pub struct Task { 28 | _thread: JoinHandle<()>, 29 | } 30 | impl Task { 31 | pub fn new(assyst: ThreadSafeAssyst, interval: Duration, callback: TaskRun) -> Task { 32 | let thread = spawn(async move { 33 | loop { 34 | callback(assyst.clone()).await; 35 | sleep(interval).await; 36 | } 37 | }); 38 | 39 | Task { _thread: thread } 40 | } 41 | 42 | pub fn new_delayed(assyst: ThreadSafeAssyst, interval: Duration, delay: Duration, callback: TaskRun) -> Task { 43 | let thread = spawn(async move { 44 | sleep(delay).await; 45 | loop { 46 | callback(assyst.clone()).await; 47 | sleep(interval).await; 48 | } 49 | }); 50 | 51 | Task { _thread: thread } 52 | } 53 | } 54 | -------------------------------------------------------------------------------- /assyst-core/src/task/tasks/get_premium_users.rs: -------------------------------------------------------------------------------- 1 | use assyst_common::config::CONFIG; 2 | use assyst_common::err; 3 | use tracing::info; 4 | 5 | use crate::assyst::ThreadSafeAssyst; 6 | use crate::rest::patreon::{Patron, PatronTier}; 7 | 8 | /// Synchronises Assyst with an updated list of patrons. 9 | pub async fn get_premium_users(assyst: ThreadSafeAssyst) { 10 | let mut premium_users: Vec = vec![]; 11 | 12 | if !CONFIG.dev.disable_patreon_synchronisation { 13 | info!("Synchronising patron list"); 14 | 15 | // get patron list and update in assyst 16 | let patrons = match crate::rest::patreon::get_patrons(&assyst.reqwest_client).await { 17 | Ok(p) => p, 18 | Err(e) => { 19 | err!("Failed to get patron list for synchronisation: {:#}", e.to_string()); 20 | return; 21 | }, 22 | }; 23 | 24 | premium_users.extend(patrons.into_iter()); 25 | 26 | info!("Synchronised patrons from Patreon: total {}", premium_users.len()); 27 | } 28 | 29 | // todo: load premium users via entitlements once twilight supports this 30 | 31 | for i in &CONFIG.dev.admin_users { 32 | premium_users.push(Patron { 33 | user_id: *i, 34 | tier: PatronTier::Tier4, 35 | _admin: true, 36 | }); 37 | } 38 | 39 | assyst.update_premium_user_list(premium_users.clone()); 40 | } 41 | -------------------------------------------------------------------------------- /assyst-core/src/task/tasks/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod get_premium_users; 2 | pub mod refresh_entitlements; 3 | pub mod reminders; 4 | pub mod top_gg_stats; 5 | -------------------------------------------------------------------------------- /assyst-core/src/task/tasks/refresh_entitlements.rs: -------------------------------------------------------------------------------- 1 | use std::collections::HashMap; 2 | 3 | use assyst_common::err; 4 | use assyst_common::macros::handle_log; 5 | use assyst_database::model::active_guild_premium_entitlement::ActiveGuildPremiumEntitlement; 6 | use tracing::info; 7 | 8 | use crate::assyst::ThreadSafeAssyst; 9 | 10 | pub async fn refresh_entitlements(assyst: ThreadSafeAssyst) { 11 | let additional = match assyst.http_client.entitlements(assyst.application_id).await { 12 | Ok(x) => match x.model().await { 13 | Ok(e) => e, 14 | Err(e) => { 15 | err!("Failed to get potential new entitlements: {e:?}"); 16 | vec![] 17 | }, 18 | }, 19 | Err(e) => { 20 | err!("Failed to get potential new entitlements: {e:?}"); 21 | vec![] 22 | }, 23 | }; 24 | 25 | for a in additional.clone() { 26 | if !assyst.entitlements.lock().unwrap().contains_key(&(a.id.get() as i64)) { 27 | let active = match ActiveGuildPremiumEntitlement::try_from(a) { 28 | Ok(a) => a, 29 | Err(e) => { 30 | err!("Error processing new entitlement: {e:?}"); 31 | continue; 32 | }, 33 | }; 34 | 35 | if active.expired() { 36 | break; 37 | } 38 | 39 | if let Err(e) = active.set(&assyst.database_handler).await { 40 | err!("Error adding new entitlement for ID {}: {e:?}", active.entitlement_id); 41 | }; 42 | handle_log(format!("New entitlement! Guild: {}", active.guild_id)); 43 | 44 | assyst 45 | .entitlements 46 | .lock() 47 | .unwrap() 48 | .insert(active.entitlement_id, active); 49 | } 50 | } 51 | 52 | let db_entitlements = ActiveGuildPremiumEntitlement::get_all(&assyst.database_handler) 53 | .await 54 | .ok() 55 | .unwrap_or(HashMap::new()); 56 | 57 | // remove entitlements from the db that are not in the rest response 58 | for entitlement in db_entitlements.values() { 59 | if !additional 60 | .iter() 61 | .any(|x| x.id.get() as i64 == entitlement.entitlement_id) 62 | || entitlement.expired() 63 | { 64 | assyst.entitlements.lock().unwrap().remove(&entitlement.entitlement_id); 65 | info!( 66 | "Removed expired entitlement {} (guild {})", 67 | entitlement.entitlement_id, entitlement.guild_id 68 | ); 69 | if let Err(e) = 70 | ActiveGuildPremiumEntitlement::delete(&assyst.database_handler, entitlement.entitlement_id).await 71 | { 72 | err!( 73 | "Error deleting existing entitlement {}: {e:?}", 74 | entitlement.entitlement_id 75 | ); 76 | } 77 | } 78 | } 79 | } 80 | -------------------------------------------------------------------------------- /assyst-core/src/task/tasks/reminders.rs: -------------------------------------------------------------------------------- 1 | use assyst_common::err; 2 | use assyst_common::util::discord::{dm_message_link, message_link}; 3 | use assyst_database::model::reminder::Reminder; 4 | use twilight_model::channel::message::AllowedMentions; 5 | use twilight_model::id::marker::{ChannelMarker, UserMarker}; 6 | use twilight_model::id::Id; 7 | 8 | use crate::assyst::ThreadSafeAssyst; 9 | 10 | // 30 seconds 11 | pub static FETCH_INTERVAL: i64 = 30000; 12 | 13 | async fn process_single_reminder(assyst: ThreadSafeAssyst, reminder: &Reminder) -> anyhow::Result<()> { 14 | assyst 15 | .http_client 16 | .create_message(Id::::new(reminder.channel_id as u64)) 17 | .allowed_mentions(Some(&AllowedMentions { 18 | parse: vec![], 19 | replied_user: false, 20 | roles: vec![], 21 | users: vec![Id::::new(reminder.user_id as u64)], 22 | })) 23 | .content(&format!( 24 | "<@{}> Reminder: {}\n{}", 25 | reminder.user_id, 26 | reminder.message, 27 | // may not be set in a guild or have a message id 28 | if reminder.guild_id != 0 && reminder.message_id != 0 { 29 | message_link( 30 | reminder.guild_id as u64, 31 | reminder.channel_id as u64, 32 | reminder.message_id as u64, 33 | ) 34 | } else if reminder.guild_id == 0 && reminder.message_id != 0 { 35 | dm_message_link(reminder.channel_id as u64, reminder.message_id as u64) 36 | } else { 37 | String::new() 38 | } 39 | )) 40 | .await?; 41 | 42 | Ok(()) 43 | } 44 | 45 | async fn process_reminders(assyst: ThreadSafeAssyst, reminders: Vec) -> Result<(), anyhow::Error> { 46 | if reminders.is_empty() { 47 | return Ok(()); 48 | } 49 | 50 | for reminder in &reminders { 51 | if let Err(e) = process_single_reminder(assyst.clone(), reminder).await { 52 | err!("Failed to process reminder: {:?}", e); 53 | } 54 | 55 | // Once we're done, delete them from database 56 | reminder.remove(&assyst.database_handler).await?; 57 | } 58 | 59 | Ok(()) 60 | } 61 | 62 | pub async fn handle_reminders(assyst: ThreadSafeAssyst) { 63 | let reminders = Reminder::fetch_expiring_max(&assyst.database_handler, FETCH_INTERVAL).await; 64 | 65 | match reminders { 66 | Ok(reminders) => { 67 | if let Err(e) = process_reminders(assyst.clone(), reminders).await { 68 | err!("Processing reminder queue failed: {:?}", e); 69 | } 70 | }, 71 | Err(e) => { 72 | err!("Fetching reminders failed: {:?}", e); 73 | }, 74 | } 75 | } 76 | -------------------------------------------------------------------------------- /assyst-core/src/task/tasks/top_gg_stats.rs: -------------------------------------------------------------------------------- 1 | use assyst_common::err; 2 | use tracing::debug; 3 | 4 | use crate::assyst::ThreadSafeAssyst; 5 | use crate::rest::top_gg::post_top_gg_stats as post_stats; 6 | 7 | pub async fn post_top_gg_stats(assyst: ThreadSafeAssyst) { 8 | debug!("Updating stats on top.gg"); 9 | 10 | if let Err(e) = post_stats( 11 | &assyst.reqwest_client, 12 | assyst.metrics_handler.guilds.with_label_values(&["guilds"]).get() as u64, 13 | ) 14 | .await 15 | { 16 | err!("Failed to post top.gg stats: {}", e.to_string()); 17 | } 18 | 19 | debug!("Updated stats on top.gg"); 20 | } 21 | -------------------------------------------------------------------------------- /assyst-database/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "assyst-database" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | futures = "0.3.12" 10 | sqlx = { version = "0.7.3", features = [ 11 | "postgres", 12 | "runtime-tokio-native-tls", 13 | ] } 14 | tokio = { workspace = true } 15 | twilight-model = { workspace = true } 16 | serde = { workspace = true } 17 | moka = { version = "0.12.3", features = ["sync"] } 18 | anyhow = { workspace = true } 19 | tracing = { workspace = true } 20 | 21 | [lints] 22 | workspace = true 23 | -------------------------------------------------------------------------------- /assyst-database/README.md: -------------------------------------------------------------------------------- 1 | # assyst-database 2 | 3 | Main database handling crate for Assyst. Handles all interactions with the database, in this case [PostgreSQL](https://www.postgresql.org/), and also contains abstrated caching logic for frequently accessed areas of the database (for example, prefixes on message commands). 4 | 5 | This crate is split into multiple separate structs, each one responsible for a table. In addition, the `impl` of each struct contains reading and writing methods for easy interfacing with that table, allowing the storage and retieval of data without needing to worry about the SQL queries involved. The function of each struct is documented using doc comments on the struct itself. -------------------------------------------------------------------------------- /assyst-database/src/lib.rs: -------------------------------------------------------------------------------- 1 | #![feature(trait_alias)] 2 | 3 | use std::borrow::Cow; 4 | 5 | use cache::DatabaseCache; 6 | use sqlx::postgres::{PgPool, PgPoolOptions}; 7 | use tracing::info; 8 | 9 | mod cache; 10 | pub mod model; 11 | 12 | static MAX_CONNECTIONS: u32 = 1; 13 | 14 | #[derive(sqlx::FromRow, Debug)] 15 | pub struct Count { 16 | pub count: i64, 17 | } 18 | 19 | #[derive(sqlx::FromRow, Debug)] 20 | pub struct DatabaseSize { 21 | pub size: String, 22 | } 23 | 24 | /// Database hendler providing a connection to the database and helper methods for inserting, 25 | /// fetching, deleting and modifying Assyst database data. 26 | pub struct DatabaseHandler { 27 | pool: PgPool, 28 | pub cache: DatabaseCache, 29 | } 30 | impl DatabaseHandler { 31 | pub async fn new(url: String, safe_url: String) -> anyhow::Result { 32 | info!( 33 | "Connecting to database on {} with {} max connections", 34 | safe_url, MAX_CONNECTIONS 35 | ); 36 | 37 | let pool = PgPoolOptions::new() 38 | .max_connections(MAX_CONNECTIONS) 39 | .connect(&url) 40 | .await?; 41 | 42 | info!("Connected to database on {}", safe_url); 43 | let cache = DatabaseCache::new(); 44 | Ok(Self { pool, cache }) 45 | } 46 | 47 | pub async fn database_size(&self) -> anyhow::Result { 48 | let query = r"SELECT pg_size_pretty(pg_database_size('assyst')) as size"; 49 | 50 | Ok(sqlx::query_as::<_, DatabaseSize>(query).fetch_one(&self.pool).await?) 51 | } 52 | } 53 | 54 | pub(crate) fn is_unique_violation(error: &sqlx::Error) -> bool { 55 | const UNIQUE_CONSTRAINT_VIOLATION_CODE: Cow<'_, str> = Cow::Borrowed("23505"); 56 | error.as_database_error().and_then(sqlx::error::DatabaseError::code) == Some(UNIQUE_CONSTRAINT_VIOLATION_CODE) 57 | } 58 | -------------------------------------------------------------------------------- /assyst-database/src/model/active_guild_premium_entitlement.rs: -------------------------------------------------------------------------------- 1 | use std::collections::HashMap; 2 | use std::time::{SystemTime, UNIX_EPOCH}; 3 | 4 | use anyhow::bail; 5 | use twilight_model::application::monetization::Entitlement; 6 | use twilight_model::util::Timestamp; 7 | 8 | use crate::{is_unique_violation, DatabaseHandler}; 9 | 10 | #[derive(sqlx::FromRow, Clone, Debug)] 11 | pub struct ActiveGuildPremiumEntitlement { 12 | pub entitlement_id: i64, 13 | pub guild_id: i64, 14 | pub user_id: i64, 15 | pub started_unix_ms: i64, 16 | pub expiry_unix_ms: i64, 17 | } 18 | impl ActiveGuildPremiumEntitlement { 19 | pub async fn set(&self, handler: &DatabaseHandler) -> anyhow::Result { 20 | let query = r"INSERT INTO active_guild_premium_entitlements VALUES ($1, $2, $3, $4, $5)"; 21 | 22 | Ok(sqlx::query(query) 23 | .bind(self.entitlement_id) 24 | .bind(self.guild_id) 25 | .bind(self.user_id) 26 | .bind(self.started_unix_ms) 27 | .bind(self.expiry_unix_ms) 28 | .execute(&handler.pool) 29 | .await 30 | .map(|_| true) 31 | .or_else(|e| if is_unique_violation(&e) { Ok(false) } else { Err(e) })?) 32 | } 33 | 34 | pub async fn delete(handler: &DatabaseHandler, entitlement_id: i64) -> anyhow::Result<()> { 35 | let query = r"DELETE FROM active_guild_premium_entitlements WHERE entitlement_id = $1"; 36 | sqlx::query(query).bind(entitlement_id).execute(&handler.pool).await?; 37 | 38 | Ok(()) 39 | } 40 | 41 | /// Useful on `ENTITLEMENT_UPDATE` where the user got billed and the expiry changes 42 | pub async fn update(&self, handler: &DatabaseHandler) -> anyhow::Result { 43 | let query = r"UPDATE active_guild_premium_entitlements SET guild_id = $2, user_id = $3, started_unix_ms = $4, expiry_unix_ms = $5 WHERE entitlement_id = $1"; 44 | 45 | Ok(sqlx::query(query) 46 | .bind(self.entitlement_id) 47 | .bind(self.guild_id) 48 | .bind(self.user_id) 49 | .bind(self.started_unix_ms) 50 | .bind(self.expiry_unix_ms) 51 | .execute(&handler.pool) 52 | .await 53 | .map(|_| true) 54 | .or_else(|e| if is_unique_violation(&e) { Ok(false) } else { Err(e) })?) 55 | } 56 | 57 | pub async fn get_all(handler: &DatabaseHandler) -> anyhow::Result> { 58 | let query = "SELECT * FROM active_guild_premium_entitlements"; 59 | let rows = sqlx::query_as::<_, Self>(query).fetch_all(&handler.pool).await?; 60 | let mut out = HashMap::new(); 61 | for r in rows { 62 | out.insert(r.entitlement_id, r); 63 | } 64 | 65 | Ok(out) 66 | } 67 | 68 | #[must_use] pub fn expired(&self) -> bool { 69 | let current = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_millis(); 70 | current > self.expiry_unix_ms as u128 && self.expiry_unix_ms != 0 71 | } 72 | } 73 | impl TryFrom for ActiveGuildPremiumEntitlement { 74 | type Error = anyhow::Error; 75 | 76 | fn try_from(value: Entitlement) -> Result { 77 | let Some(guild_id) = value.guild_id else { 78 | bail!( 79 | "Entitlement ID {} (guild {:?} user {:?}) has no associated guild!", 80 | value.id, 81 | value.guild_id, 82 | value.user_id 83 | ) 84 | }; 85 | 86 | let Some(user_id) = value.user_id else { 87 | bail!( 88 | "Entitlement ID {} (guild {:?} user {:?}) has no associated user!", 89 | value.id, 90 | value.guild_id, 91 | value.user_id 92 | ) 93 | }; 94 | 95 | // no expiry/created = test entitlement, requires special handling 96 | let active = Self { 97 | entitlement_id: value.id.get() as i64, 98 | guild_id: guild_id.get() as i64, 99 | user_id: user_id.get() as i64, 100 | started_unix_ms: value 101 | .starts_at 102 | .unwrap_or(Timestamp::from_micros(0).unwrap()) 103 | .as_micros() 104 | / 1000, 105 | expiry_unix_ms: value.ends_at.unwrap_or(Timestamp::from_micros(0).unwrap()).as_micros() / 1000, 106 | }; 107 | 108 | Ok(active) 109 | } 110 | } 111 | -------------------------------------------------------------------------------- /assyst-database/src/model/badtranslator_channel.rs: -------------------------------------------------------------------------------- 1 | use crate::{is_unique_violation, DatabaseHandler}; 2 | 3 | #[derive(sqlx::FromRow)] 4 | pub struct BadTranslatorChannel { 5 | pub id: i64, 6 | pub target_language: String, 7 | } 8 | impl BadTranslatorChannel { 9 | pub async fn get_all(handler: &DatabaseHandler) -> anyhow::Result> { 10 | let query = "SELECT * FROM bt_channels"; 11 | let rows = sqlx::query_as::<_, Self>(query).fetch_all(&handler.pool).await?; 12 | Ok(rows) 13 | } 14 | 15 | pub async fn delete(handler: &DatabaseHandler, id: i64) -> anyhow::Result { 16 | let query = r"DELETE FROM bt_channels WHERE id = $1 RETURNING *"; 17 | 18 | Ok(sqlx::query(query) 19 | .bind(id) 20 | .fetch_all(&handler.pool) 21 | .await 22 | .map(|x| !x.is_empty())?) 23 | } 24 | 25 | pub async fn update_language(&self, handler: &DatabaseHandler, new_language: &str) -> anyhow::Result { 26 | let query = r"UPDATE bt_channels SET target_language = $1 WHERE id = $2 RETURNING *"; 27 | 28 | Ok(sqlx::query(query) 29 | .bind(new_language) 30 | .bind(self.id) 31 | .fetch_all(&handler.pool) 32 | .await 33 | .map(|x| !x.is_empty())?) 34 | } 35 | 36 | pub async fn set(&self, handler: &DatabaseHandler) -> anyhow::Result { 37 | let query = r"INSERT INTO bt_channels VALUES ($1, $2)"; 38 | 39 | sqlx::query(query) 40 | .bind(self.id) 41 | .bind(&self.target_language) 42 | .execute(&handler.pool) 43 | .await 44 | .map(|_| true) 45 | .or_else(|e| { 46 | if is_unique_violation(&e) { 47 | Ok(false) 48 | } else { 49 | Err(e.into()) 50 | } 51 | }) 52 | } 53 | } 54 | -------------------------------------------------------------------------------- /assyst-database/src/model/badtranslator_messages.rs: -------------------------------------------------------------------------------- 1 | use crate::DatabaseHandler; 2 | 3 | pub struct BadTranslatorMessages { 4 | _guild_id: i64, 5 | _message_count: i64, 6 | } 7 | impl BadTranslatorMessages { 8 | pub async fn increment(handler: &DatabaseHandler, guild_id: i64) -> anyhow::Result<()> { 9 | let query = "insert into bt_messages (guild_id, message_count) values ($1, 1) on conflict (guild_id) do update set message_count = bt_messages.message_count + 1 where bt_messages.guild_id = $1;"; 10 | sqlx::query(query).bind(guild_id).execute(&handler.pool).await?; 11 | Ok(()) 12 | } 13 | } 14 | -------------------------------------------------------------------------------- /assyst-database/src/model/colour_role.rs: -------------------------------------------------------------------------------- 1 | use crate::DatabaseHandler; 2 | 3 | /// A colour role is a self-assignable role to grant a colour to a user. 4 | #[derive(sqlx::FromRow, Debug, Clone)] 5 | pub struct ColourRole { 6 | pub role_id: i64, 7 | pub name: String, 8 | pub guild_id: i64, 9 | } 10 | impl ColourRole { 11 | /// List all colour roles in a guild. 12 | pub async fn list_in_guild(handler: &DatabaseHandler, guild_id: i64) -> Result, sqlx::Error> { 13 | if let Some(c) = handler.cache.get_guild_colour_roles(guild_id as u64) { 14 | return Ok(c); 15 | } 16 | 17 | let query = r"SELECT * FROM colors WHERE guild_id = $1"; 18 | 19 | let roles: Vec = sqlx::query_as(query).bind(guild_id).fetch_all(&handler.pool).await?; 20 | handler.cache.insert_guild_colour_roles(guild_id as u64, roles.clone()); 21 | 22 | Ok(roles) 23 | } 24 | 25 | /// Inser a new colour role. 26 | pub async fn insert(&self, handler: &DatabaseHandler) -> Result<(), sqlx::Error> { 27 | let query = r"INSERT INTO colors VALUES ($1, $2, $3)"; 28 | 29 | sqlx::query(query) 30 | .bind(self.role_id) 31 | .bind(self.name.clone()) 32 | .bind(self.guild_id) 33 | .execute(&handler.pool) 34 | .await 35 | .map(|_| ()) 36 | } 37 | 38 | /// Remove a colour role. Returns true on successful removal, false if the role did not exist. 39 | pub async fn remove(&self, handler: &DatabaseHandler) -> Result { 40 | let query = r"DELETE FROM colors WHERE guild_id = $1 AND name = $2 RETURNING *"; 41 | 42 | let result = sqlx::query_as::<_, ColourRole>(query) 43 | .bind(self.guild_id) 44 | .bind(self.name.clone()) 45 | .fetch_one(&handler.pool) 46 | .await; 47 | 48 | match result { 49 | Ok(_) => Ok(true), 50 | Err(sqlx::Error::RowNotFound) => Ok(false), 51 | Err(e) => Err(e), 52 | } 53 | } 54 | } 55 | -------------------------------------------------------------------------------- /assyst-database/src/model/command_usage.rs: -------------------------------------------------------------------------------- 1 | use crate::DatabaseHandler; 2 | 3 | #[derive(sqlx::FromRow)] 4 | pub struct CommandUsage { 5 | pub command_name: String, 6 | pub uses: i32, 7 | } 8 | impl CommandUsage { 9 | pub async fn get_command_usage_stats(handler: &DatabaseHandler) -> Result, sqlx::Error> { 10 | let query = "SELECT * FROM command_uses order by uses desc"; 11 | sqlx::query_as::<_, Self>(query).fetch_all(&handler.pool).await 12 | } 13 | 14 | pub async fn get_command_usage_stats_for(&self, handler: &DatabaseHandler) -> Result { 15 | let query = "SELECT * FROM command_uses where command_name = $1 order by uses desc"; 16 | sqlx::query_as::<_, Self>(query) 17 | .bind(&self.command_name) 18 | .fetch_one(&handler.pool) 19 | .await 20 | } 21 | 22 | pub async fn increment_command_uses(&self, handler: &DatabaseHandler) -> Result<(), sqlx::Error> { 23 | let query = "insert into command_uses (command_name, uses) values ($1, 1) on conflict (command_name) do update set uses = command_uses.uses + 1 where command_uses.command_name = $1;"; 24 | sqlx::query(query) 25 | .bind(&self.command_name) 26 | .execute(&handler.pool) 27 | .await?; 28 | Ok(()) 29 | } 30 | } 31 | -------------------------------------------------------------------------------- /assyst-database/src/model/free_tier_2_requests.rs: -------------------------------------------------------------------------------- 1 | use crate::DatabaseHandler; 2 | 3 | #[derive(sqlx::FromRow, Debug)] 4 | pub struct FreeTier2Requests { 5 | pub user_id: i64, 6 | pub count: i32, 7 | } 8 | impl FreeTier2Requests { 9 | #[must_use] pub fn new(user_id: u64) -> FreeTier2Requests { 10 | FreeTier2Requests { 11 | user_id: user_id as i64, 12 | count: 0, 13 | } 14 | } 15 | 16 | pub async fn change_free_tier_2_requests( 17 | &self, 18 | handler: &DatabaseHandler, 19 | change_amount: i64, 20 | ) -> anyhow::Result<()> { 21 | let query = "insert into free_tier1_requests values($1, $2) on conflict (user_id) do update set count = free_tier1_requests.count + $2 where free_tier1_requests.user_id = $1"; 22 | 23 | sqlx::query(query) 24 | .bind(self.user_id) 25 | .bind(change_amount) 26 | .execute(&handler.pool) 27 | .await?; 28 | 29 | Ok(()) 30 | } 31 | 32 | pub async fn get_user_free_tier_2_requests( 33 | handler: &DatabaseHandler, 34 | user_id: u64, 35 | ) -> anyhow::Result { 36 | let fetch_query = "select * from free_tier1_requests where user_id = $1"; 37 | 38 | match sqlx::query_as::<_, FreeTier2Requests>(fetch_query) 39 | .bind(user_id as i64) 40 | .fetch_one(&handler.pool) 41 | .await 42 | { 43 | Ok(x) => Ok(x), 44 | Err(sqlx::Error::RowNotFound) => Ok(FreeTier2Requests { 45 | user_id: user_id as i64, 46 | count: 0, 47 | }), 48 | Err(e) => Err(e.into()), 49 | } 50 | } 51 | } 52 | -------------------------------------------------------------------------------- /assyst-database/src/model/global_blacklist.rs: -------------------------------------------------------------------------------- 1 | use crate::DatabaseHandler; 2 | 3 | /// The global blacklist is a list of users who are completely blacklisted from using any of the 4 | /// bot's functionality. It is a simple list table with one column of user IDs which are 5 | /// blacklisted. 6 | pub struct GlobalBlacklist {} 7 | impl GlobalBlacklist { 8 | pub async fn is_blacklisted(handler: &DatabaseHandler, user_id: u64) -> anyhow::Result { 9 | if let Some(blacklisted) = handler.cache.get_user_global_blacklist(user_id) { 10 | return Ok(blacklisted); 11 | } 12 | 13 | let query = r"SELECT user_id FROM blacklist WHERE user_id = $1"; 14 | 15 | match sqlx::query_as::<_, (i64,)>(query) 16 | .bind(user_id as i64) 17 | .fetch_one(&handler.pool) 18 | .await 19 | .map(|result| result.0) 20 | { 21 | Ok(_) => { 22 | handler.cache.set_user_global_blacklist(user_id, true); 23 | Ok(true) 24 | }, 25 | Err(sqlx::Error::RowNotFound) => { 26 | handler.cache.set_user_global_blacklist(user_id, false); 27 | Ok(false) 28 | }, 29 | Err(err) => Err(err.into()), 30 | } 31 | } 32 | 33 | pub async fn set_user_blacklisted(&self, handler: &DatabaseHandler, user_id: u64) -> anyhow::Result<()> { 34 | let query = r"INSERT INTO blacklist VALUES ($1)"; 35 | 36 | sqlx::query(query).bind(user_id as i64).execute(&handler.pool).await?; 37 | handler.cache.set_user_global_blacklist(user_id, true); 38 | 39 | Ok(()) 40 | } 41 | 42 | pub async fn remove_user_from_blacklist(&self, handler: &DatabaseHandler, user_id: u64) -> Result<(), sqlx::Error> { 43 | let query = r"DELETE FROM blacklist WHERE user_id = $1"; 44 | 45 | handler.cache.set_user_global_blacklist(user_id, false); 46 | sqlx::query(query) 47 | .bind(user_id as i64) 48 | .execute(&handler.pool) 49 | .await 50 | .map(|_| ()) 51 | } 52 | } 53 | -------------------------------------------------------------------------------- /assyst-database/src/model/guild_disabled_command.rs: -------------------------------------------------------------------------------- 1 | use anyhow::Context; 2 | 3 | use crate::DatabaseHandler; 4 | 5 | #[derive(sqlx::FromRow)] 6 | pub struct GuildDisabledCommand { 7 | pub guild_id: i64, 8 | pub command_name: String, 9 | } 10 | impl GuildDisabledCommand { 11 | pub async fn is_disabled(&self, handler: &DatabaseHandler) -> anyhow::Result { 12 | if let Some(commands) = handler.cache.get_guild_disabled_commands(self.guild_id as u64) { 13 | return Ok(commands.lock().unwrap().contains(&self.command_name)); 14 | } 15 | 16 | let query = "select * from disabled_commands where guild_id = $1"; 17 | let result = sqlx::query_as::<_, Self>(query) 18 | .bind(self.guild_id) 19 | .fetch_all(&handler.pool) 20 | .await 21 | .context("Failed to fetch guild disabled commands from database")?; 22 | 23 | handler.cache.reset_disabled_commands_for(self.guild_id as u64); 24 | 25 | for command in &result { 26 | handler 27 | .cache 28 | .set_command_disabled(self.guild_id as u64, &command.command_name); 29 | } 30 | 31 | let is_disabled = result.iter().any(|cmd| cmd.command_name == self.command_name); 32 | Ok(is_disabled) 33 | } 34 | 35 | pub async fn enable(&self, handler: &DatabaseHandler) -> anyhow::Result<()> { 36 | let query = "delete from disabled_commands where guild_id = $1 and command_name = $2"; 37 | 38 | sqlx::query(query) 39 | .bind(self.guild_id) 40 | .bind(&self.command_name) 41 | .execute(&handler.pool) 42 | .await?; 43 | 44 | handler 45 | .cache 46 | .set_command_enabled(self.guild_id as u64, &self.command_name); 47 | 48 | Ok(()) 49 | } 50 | 51 | pub async fn disable(&self, handler: &DatabaseHandler) -> anyhow::Result<()> { 52 | let query = "insert into disabled_commands(guild_id, command_name) values($1, $2)"; 53 | 54 | sqlx::query(query) 55 | .bind(self.guild_id) 56 | .bind(&self.command_name) 57 | .execute(&handler.pool) 58 | .await?; 59 | 60 | handler 61 | .cache 62 | .set_command_disabled(self.guild_id as u64, &self.command_name); 63 | 64 | Ok(()) 65 | } 66 | } 67 | -------------------------------------------------------------------------------- /assyst-database/src/model/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod active_guild_premium_entitlement; 2 | pub mod badtranslator_channel; 3 | pub mod badtranslator_messages; 4 | pub mod colour_role; 5 | pub mod command_usage; 6 | pub mod free_tier_2_requests; 7 | pub mod global_blacklist; 8 | pub mod guild_disabled_command; 9 | pub mod prefix; 10 | pub mod reminder; 11 | pub mod tag; 12 | pub mod user_votes; 13 | -------------------------------------------------------------------------------- /assyst-database/src/model/prefix.rs: -------------------------------------------------------------------------------- 1 | use std::hash::Hash; 2 | 3 | use crate::DatabaseHandler; 4 | 5 | #[derive(Clone, Hash, PartialEq)] 6 | /// A Prefix is a unique identifier for invocating message-based commands in a guild. This table is 7 | /// relatively simple, holding only a guild ID and its associated prefix, since each guild can only 8 | /// have one prefix at a time. 9 | pub struct Prefix { 10 | pub prefix: String, 11 | } 12 | impl Prefix { 13 | pub async fn set(&self, handler: &DatabaseHandler, guild_id: u64) -> anyhow::Result<()> { 14 | let query = r"INSERT INTO prefixes(guild, prefix) VALUES($1, $2) ON CONFLICT (guild) DO UPDATE SET prefix = $2 WHERE prefixes.guild = $1"; 15 | 16 | sqlx::query(query) 17 | .bind(guild_id as i64) 18 | .bind(self.clone().prefix) 19 | .execute(&handler.pool) 20 | .await?; 21 | 22 | handler.cache.set_prefix(guild_id, self.clone()); 23 | 24 | Ok(()) 25 | } 26 | 27 | pub async fn get(handler: &DatabaseHandler, guild_id: u64) -> anyhow::Result> { 28 | if let Some(prefix) = handler.cache.get_prefix(guild_id) { 29 | return Ok(Some(prefix)); 30 | } 31 | 32 | let query = "SELECT * FROM prefixes WHERE guild = $1"; 33 | 34 | match sqlx::query_as::<_, (String,)>(query) 35 | .bind(guild_id as i64) 36 | .fetch_one(&handler.pool) 37 | .await 38 | { 39 | Ok(res) => { 40 | let prefix = Prefix { prefix: res.0 }; 41 | handler.cache.set_prefix(guild_id, prefix.clone()); 42 | Ok(Some(prefix)) 43 | }, 44 | Err(sqlx::Error::RowNotFound) => Ok(None), 45 | Err(err) => Err(err.into()), 46 | } 47 | } 48 | 49 | #[must_use] pub fn size_of(&self) -> u64 { 50 | self.prefix.as_bytes().len() as u64 51 | } 52 | } 53 | -------------------------------------------------------------------------------- /assyst-database/src/model/reminder.rs: -------------------------------------------------------------------------------- 1 | use std::time::{SystemTime, UNIX_EPOCH}; 2 | 3 | use crate::DatabaseHandler; 4 | 5 | #[derive(sqlx::FromRow, Debug)] 6 | pub struct Reminder { 7 | pub id: i32, 8 | pub user_id: i64, 9 | pub timestamp: i64, 10 | pub guild_id: i64, 11 | pub channel_id: i64, 12 | pub message_id: i64, 13 | pub message: String, 14 | } 15 | impl Reminder { 16 | /// Fetch all reminders with a certain maximum expiration 17 | pub async fn fetch_expiring_max(handler: &DatabaseHandler, time_delta: i64) -> Result, sqlx::Error> { 18 | let query = "SELECT * FROM reminders WHERE timestamp < $1"; 19 | 20 | let unix: i64 = SystemTime::now() 21 | .duration_since(UNIX_EPOCH) 22 | .expect("Time went backwards") 23 | .as_millis() 24 | .try_into() 25 | .expect("count not fit u128 into target type"); 26 | 27 | sqlx::query_as::<_, Self>(query) 28 | .bind(unix + time_delta) 29 | .fetch_all(&handler.pool) 30 | .await 31 | } 32 | 33 | /// Fetch all reminders within a certain count for a user ID 34 | pub async fn fetch_user_reminders( 35 | handler: &DatabaseHandler, 36 | user: u64, 37 | count: u64, 38 | ) -> Result, sqlx::Error> { 39 | let query = r"SELECT * FROM reminders WHERE user_id = $1 ORDER BY timestamp ASC LIMIT $2"; 40 | 41 | sqlx::query_as::<_, Self>(query) 42 | .bind(user as i64) 43 | .bind(count as i64) 44 | .fetch_all(&handler.pool) 45 | .await 46 | } 47 | 48 | /// True on successful remove, false otherwise 49 | pub async fn remove(&self, handler: &DatabaseHandler) -> Result { 50 | let query = r"DELETE FROM reminders WHERE user_id = $1 AND id = $2 RETURNING *"; 51 | 52 | sqlx::query(query) 53 | .bind(self.user_id) 54 | .bind(self.id) 55 | .fetch_all(&handler.pool) 56 | .await 57 | .map(|s| !s.is_empty()) 58 | } 59 | 60 | /// Add a new reminder 61 | pub async fn insert(&self, handler: &DatabaseHandler) -> Result<(), sqlx::Error> { 62 | let query = r"INSERT INTO reminders VALUES ($1, $2, $3, $4, $5, $6)"; 63 | 64 | sqlx::query(query) 65 | .bind(self.user_id) 66 | .bind(self.timestamp) 67 | .bind(self.guild_id) 68 | .bind(self.channel_id) 69 | .bind(self.message_id) 70 | .bind(&*self.message) 71 | .execute(&handler.pool) 72 | .await 73 | .map(|_| ()) 74 | } 75 | } 76 | -------------------------------------------------------------------------------- /assyst-database/src/model/user_votes.rs: -------------------------------------------------------------------------------- 1 | use crate::DatabaseHandler; 2 | 3 | /// A Voter is a user with a certain number of accrued votes. In the old Assyst, it was possible to 4 | /// get a leaderboard of the top voters with their username and discriminator (hence these fields), 5 | /// but these have since become unused. In a future version, they may be removed entirely. 6 | #[derive(sqlx::FromRow, Debug)] 7 | pub struct UserVotes { 8 | pub user_id: i64, 9 | pub username: String, 10 | pub discriminator: String, 11 | pub count: i32, 12 | } 13 | impl UserVotes { 14 | pub async fn get_user_votes(handler: &DatabaseHandler, user_id: u64) -> anyhow::Result> { 15 | let fetch_query = "select * from user_votes where user_id = $1"; 16 | 17 | let result = sqlx::query_as::<_, UserVotes>(fetch_query) 18 | .bind(user_id as i64) 19 | .fetch_optional(&handler.pool) 20 | .await?; 21 | 22 | Ok(result) 23 | } 24 | 25 | pub async fn increment_user_votes( 26 | handler: &DatabaseHandler, 27 | user_id: u64, 28 | username: &str, 29 | discriminator: &str, 30 | ) -> anyhow::Result<()> { 31 | let query = "insert into user_votes values($1, $2, $3, 1) on conflict (user_id) do update set count = user_votes.count + 1 where user_votes.user_id = $1"; 32 | 33 | sqlx::query(query) 34 | .bind(user_id as i64) 35 | .bind(username) 36 | .bind(discriminator) 37 | .execute(&handler.pool) 38 | .await?; 39 | 40 | Ok(()) 41 | } 42 | } 43 | -------------------------------------------------------------------------------- /assyst-flux-iface/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "assyst-flux-iface" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | [dependencies] 7 | assyst-common = { path = "../assyst-common" } 8 | assyst-database = { path = "../assyst-database" } 9 | tokio = { workspace = true } 10 | anyhow = { workspace = true } 11 | libc = "0.2.155" 12 | serde = { workspace = true } 13 | serde_json = "1.0.121" 14 | 15 | [lints] 16 | workspace = true 17 | -------------------------------------------------------------------------------- /assyst-flux-iface/README.md: -------------------------------------------------------------------------------- 1 | # assyst-flux-iface 2 | 3 | Basic wrapper over Flux. Provides utilities for performing operations, creating full command-line scripts, compiling the binary, etc. \ 4 | Uses values from assyst-common `CONFIG` for details such as the workspace path. -------------------------------------------------------------------------------- /assyst-flux-iface/src/flux_request.rs: -------------------------------------------------------------------------------- 1 | use std::collections::HashMap; 2 | 3 | use super::limits::LimitData; 4 | 5 | /// A step in a Flux execution. 6 | pub enum FluxStep { 7 | /// Input file. Saves the file and passes to Flux as `-i path`. Input must be the first step. 8 | Input(Vec), 9 | /// Operation. Passes to Flux as `-o operation[k=v]` 10 | Operation((String, HashMap)), 11 | /// Output. Passes to Flux as `path` at the end. Output must be the last step. 12 | Output, 13 | /// Frame limit of inputs. Inputs will have additional frames removed. 14 | ImagePageLimit(u64), 15 | /// Resolution limit of input. Input will shrink, preserving aspect ratio, to fit this. 16 | ResolutionLimit((u64, u64)), 17 | /// Whether to disable video decoding. 18 | VideoDecodeDisabled, 19 | /// Get media info 20 | Info, 21 | /// Get version info 22 | Version, 23 | } 24 | 25 | #[derive(Default)] 26 | pub struct FluxRequest(pub Vec); 27 | impl FluxRequest { 28 | #[must_use] pub fn new_with_input_and_limits(input: Vec, limits: &LimitData) -> Self { 29 | let mut new = Self(vec![]); 30 | new.input(input); 31 | new.limits(limits); 32 | new 33 | } 34 | 35 | #[must_use] pub fn new_basic(input: Vec, limits: &LimitData, operation: &str) -> Self { 36 | let mut new = Self(vec![]); 37 | new.input(input); 38 | new.limits(limits); 39 | new.operation(operation.to_owned(), HashMap::new()); 40 | new.output(); 41 | new 42 | } 43 | 44 | pub fn input(&mut self, input: Vec) { 45 | self.0.push(FluxStep::Input(input)); 46 | } 47 | 48 | pub fn operation(&mut self, name: String, options: HashMap) { 49 | self.0.push(FluxStep::Operation((name, options))); 50 | } 51 | 52 | pub fn output(&mut self) { 53 | self.0.push(FluxStep::Output); 54 | } 55 | 56 | pub fn limits(&mut self, limits: &LimitData) { 57 | self.0.push(FluxStep::ImagePageLimit(limits.frames)); 58 | self.0.push(FluxStep::ResolutionLimit((limits.size, limits.size))); 59 | if !limits.video_decode_enabled { 60 | self.0.push(FluxStep::VideoDecodeDisabled); 61 | } 62 | } 63 | 64 | pub fn info(&mut self) { 65 | self.0.push(FluxStep::Info); 66 | } 67 | 68 | pub fn version(&mut self) { 69 | self.0.push(FluxStep::Version); 70 | } 71 | } 72 | -------------------------------------------------------------------------------- /assyst-flux-iface/src/limits.rs: -------------------------------------------------------------------------------- 1 | use std::time::Duration; 2 | 3 | pub struct LimitData { 4 | pub time: Duration, 5 | pub size: u64, 6 | pub frames: u64, 7 | pub video_decode_enabled: bool, 8 | } 9 | 10 | pub const LIMITS_FREE: LimitData = LimitData { 11 | time: Duration::from_secs(40), 12 | size: 768, 13 | frames: 150, 14 | video_decode_enabled: false, 15 | }; 16 | 17 | pub const LIMITS_USER_TIER_1: LimitData = LimitData { 18 | time: Duration::from_secs(60), 19 | size: 1024, 20 | frames: 200, 21 | video_decode_enabled: true, 22 | }; 23 | 24 | pub const LIMITS_USER_TIER_2: LimitData = LimitData { 25 | time: Duration::from_secs(80), 26 | size: 2048, 27 | frames: 225, 28 | video_decode_enabled: true, 29 | }; 30 | 31 | pub const LIMITS_USER_TIER_3: LimitData = LimitData { 32 | time: Duration::from_secs(120), 33 | size: 4096, 34 | frames: 250, 35 | video_decode_enabled: true, 36 | }; 37 | 38 | pub const LIMITS_GUILD_TIER_1: LimitData = LimitData { 39 | time: Duration::from_secs(60), 40 | size: 1024, 41 | frames: 200, 42 | video_decode_enabled: true, 43 | }; 44 | 45 | #[must_use] pub fn premium_user_to_limits(tier: u64) -> LimitData { 46 | match tier { 47 | 0 => LIMITS_FREE, 48 | 1 => LIMITS_GUILD_TIER_1, 49 | 2 => LIMITS_USER_TIER_2, 50 | 3 => LIMITS_USER_TIER_3, 51 | _ => unreachable!(), 52 | } 53 | } 54 | -------------------------------------------------------------------------------- /assyst-gateway/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "assyst-gateway" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | twilight-gateway = { workspace = true } 10 | anyhow = { workspace = true } 11 | assyst-common = { path = "../assyst-common" } 12 | assyst-database = { path = "../assyst-database" } 13 | futures-util = "0.3.29" 14 | lazy_static = "1.4.0" 15 | tokio = { workspace = true } 16 | tracing = { workspace = true } 17 | tracing-subscriber = { version = "0.3.16", features = ["time", "env-filter"] } 18 | time = { version = "0.3.31", features = ["macros"] } 19 | twilight-model = { workspace = true } 20 | twilight-http = { workspace = true } 21 | jemallocator = "0.5.4" 22 | aws-lc-rs = "1.10.0" 23 | rustls = "0.23.15" 24 | 25 | [lints] 26 | workspace = true 27 | -------------------------------------------------------------------------------- /assyst-gateway/README.md: -------------------------------------------------------------------------------- 1 | # assyst-gateway 2 | 3 | This is the basic HTTP connection to the Discord WebSocket gateway, in order to receive messages (known as events in Discord) for processing. These messages include messages being sent, updated (such as being edited), deleted, as well as other events such as the bot joining a new guild, being removed from a guild, or a shard (an individual connection - see below) becoming ready to receive events. 4 | 5 | The communication through the Discord WebSocket gateway is fairly complex, expecially for large bots that are in many Discord guilds. The details of the low-level communication are abstracted away by the Discord libaray Assyst uses (https://crates.io/crates/twilight-gateway), but there are still some components that must be considered. 6 | 7 | One of the key considerations of the design for Assyst2 is very reliable operation, including fast restarts. The production instance of Assyst is present in many thousands of guilds, and as such can process hundreds of events per second. Usually, this is far too much for a single WebSocket connection to cope with, so communication with Discord is divided into multiple WebSocket connections known as *shards*. The issue is that shards have a certain 'startup time' and only one can start at once, meaning that for a large instance, this gateway process can take several minutes to fully connect to Discord - not ideal. 8 | 9 | As a result, this crate exists to allow updates and restarts to other parts of the Assyst infrastructure (primarily assyst-core, for command updates) without forcing a reconnection to the Discord WebSocket gateway, allowing much faster updates. 10 | 11 | ![image](./assets/gateway.png) 12 | 13 | When messages are received by the Assyst gateway, they are immediately sent via a UNIX pipe (`UnixStream` in Tokio) for parsing and processing. In addition, when the core restarts or crashes, messages are buffered for when the core reconnects. This way, no messages are lost and so in theory, updates to the core should have zero percieved downtime. -------------------------------------------------------------------------------- /assyst-gateway/assets/gateway.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jacherr/assyst2/faf20ce0919e82a8c703a2f53214ef493004f6c9/assyst-gateway/assets/gateway.png -------------------------------------------------------------------------------- /assyst-gateway/src/main.rs: -------------------------------------------------------------------------------- 1 | use assyst_common::config::CONFIG; 2 | use assyst_common::ok_or_break; 3 | use assyst_common::pipe::pipe_server::PipeServer; 4 | use assyst_common::pipe::GATEWAY_PIPE_PATH; 5 | use assyst_common::util::tracing_init; 6 | use futures_util::StreamExt; 7 | use tokio::sync::mpsc::{channel, Sender}; 8 | use tokio::{signal, spawn}; 9 | use tracing::{debug, info, trace, warn}; 10 | use twilight_gateway::{create_recommended, ConfigBuilder as GatewayConfigBuilder, Intents, Message, Shard}; 11 | use twilight_http::Client as HttpClient; 12 | use twilight_model::gateway::payload::outgoing::update_presence::UpdatePresencePayload; 13 | use twilight_model::gateway::presence::{Activity, ActivityType, Status}; 14 | 15 | // Jemallocator is probably unnecessary for the average instance, 16 | // but when handling hundreds of events per second the performance improvement 17 | // can be measurable 18 | #[cfg(target_os = "linux")] 19 | #[global_allocator] 20 | static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc; 21 | 22 | lazy_static::lazy_static! { 23 | static ref ACTIVITY: Activity = Activity { 24 | application_id: None, 25 | assets: None, 26 | created_at: None, 27 | details: None, 28 | emoji: None, 29 | flags: None, 30 | id: None, 31 | instance: None, 32 | kind: ActivityType::Playing, 33 | name: format!("{}help | jacher.io/assyst", CONFIG.prefix.default), 34 | party: None, 35 | secrets: None, 36 | state: None, 37 | timestamps: None, 38 | url: None, 39 | buttons: Vec::new(), 40 | }; 41 | } 42 | 43 | #[tokio::main] 44 | async fn main() -> anyhow::Result<()> { 45 | assert!(std::env::consts::OS == "linux", "Assyst is supported on Linux only."); 46 | 47 | rustls::crypto::aws_lc_rs::default_provider().install_default().unwrap(); 48 | 49 | tracing_init(); 50 | 51 | let http_client = HttpClient::new(CONFIG.authentication.discord_token.clone()); 52 | 53 | let presence = UpdatePresencePayload::new(vec![ACTIVITY.to_owned()], false, None, Status::Online).unwrap(); 54 | 55 | let intents = Intents::MESSAGE_CONTENT | Intents::GUILDS | Intents::GUILD_MESSAGES | Intents::DIRECT_MESSAGES; 56 | debug!("intents={:?}", intents); 57 | let gateway_config = GatewayConfigBuilder::new(CONFIG.authentication.discord_token.clone(), intents) 58 | .presence(presence) 59 | .build(); 60 | 61 | let shards = create_recommended(&http_client, gateway_config.clone(), |_, _| gateway_config.clone()) 62 | .await 63 | .unwrap() 64 | .collect::>(); 65 | 66 | info!("Recommended shard count: {}", shards.len()); 67 | 68 | // pipe thread tx/rx 69 | let (tx, mut rx) = channel::(25); 70 | 71 | let mut core_pipe_server = PipeServer::listen(GATEWAY_PIPE_PATH).unwrap(); 72 | info!("Core listener started on {}", GATEWAY_PIPE_PATH); 73 | 74 | // pipe thread 75 | tokio::spawn(async move { 76 | info!("Awaiting connection from assyst-core"); 77 | loop { 78 | if let Ok(mut stream) = core_pipe_server.accept_connection().await { 79 | info!("Connection received from assyst-core"); 80 | while let Some(v) = rx.recv().await { 81 | ok_or_break!(stream.write_string(v).await); 82 | } 83 | warn!("Connection to assyst-core lost, awaiting reconnection"); 84 | } 85 | } 86 | }); 87 | 88 | let mut tasks = vec![]; 89 | let shards_count = shards.len(); 90 | 91 | for shard in shards { 92 | info!( 93 | "Registering runner for shard {} of {}", 94 | shard.id().number(), 95 | shards_count - 1 96 | ); 97 | tasks.push(spawn(runner(shard, tx.clone()))); 98 | } 99 | 100 | signal::ctrl_c().await?; 101 | 102 | Ok(()) 103 | } 104 | 105 | async fn runner(mut shard: Shard, tx: Sender) { 106 | loop { 107 | match shard.next().await { 108 | Some(Ok(Message::Text(message))) => { 109 | trace!("got message: {message}"); 110 | let _ = tx.try_send(message); 111 | }, 112 | Some(Err(e)) => { 113 | warn!(?e, "error receiving event"); 114 | }, 115 | _ => {}, 116 | } 117 | } 118 | } 119 | -------------------------------------------------------------------------------- /assyst-proc-macro/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "assyst-proc-macro" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | [lib] 7 | proc-macro = true 8 | 9 | [dependencies] 10 | syn = "2.0.48" 11 | quote = "1.0.35" 12 | proc-macro2 = "1.0.78" 13 | 14 | [lints] 15 | workspace = true 16 | -------------------------------------------------------------------------------- /assyst-proc-macro/README.md: -------------------------------------------------------------------------------- 1 | # assyst-proc-macro 2 | 3 | This crate defines proc macros, used primarily by the `assyst-core` crate for its `#[command]` macro. 4 | 5 | NOTE: if you're looking to implement a **command group**, use the `assyst_core::command::define_commandgroup!` macro. 6 | 7 | ### `#[command]` 8 | This macro allows you to write bot commands (independent of the source: gateway or interaction) as regular async functions. 9 | 10 | The first parameter must always be the `CommandCtxt`, followed by any number of argument types. 11 | `Rest` must be the last argument, if present. For example, the `-remind 1h do something` can be defined as: 12 | ```rs 13 | #[command( 14 | name = "remind", 15 | aliases = ["reminder"], 16 | description = "get reminders or set a reminder, time format is xdyhzm (check examples)", 17 | access = Availability::Public, 18 | cooldown = Duration::from_secs(2) 19 | )] 20 | pub fn remind(ctxt: CommandCtxt<'_>, when: Time, text: Rest) -> anyhow::Result<()> { 21 | // ... 22 | Ok(()) 23 | } 24 | ``` 25 | The macro takes some metadata about the command (description, availablity, etc.) and generates a struct that implements the `Command` trait, in which it calls `::parse()` on each of the provided types (`Time`, `Rest`) and finally passes it to the annotated function. 26 | 27 | For even more details (e.g. its exact expansion), take a look at the code. There's documentation on the proc macro function, too. 28 | -------------------------------------------------------------------------------- /assyst-string-fmt/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "assyst-string-fmt" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | [dependencies] 7 | regex = "1.10.5" 8 | 9 | [lints] 10 | workspace = true 11 | -------------------------------------------------------------------------------- /assyst-string-fmt/README.md: -------------------------------------------------------------------------------- 1 | # assyst-string-fmt 2 | 3 | String formatting and parsing utilities, for example Discord markdown and ANSI colours. -------------------------------------------------------------------------------- /assyst-string-fmt/src/lib.rs: -------------------------------------------------------------------------------- 1 | pub mod ansi; 2 | pub mod markdown; 3 | 4 | pub use ansi::Ansi; 5 | pub use markdown::Markdown; 6 | -------------------------------------------------------------------------------- /assyst-tag/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "assyst-tag" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | rand = "0.8" 10 | anyhow = { workspace = true } 11 | assyst-common = { path = "../assyst-common" } 12 | assyst-string-fmt = { path = "../assyst-string-fmt" } 13 | bytes = "1.0.1" 14 | either = "1.9.0" 15 | memchr = "2.6.4" 16 | 17 | [lints] 18 | workspace = true 19 | -------------------------------------------------------------------------------- /assyst-tag/fuzz/.gitignore: -------------------------------------------------------------------------------- 1 | target 2 | corpus 3 | artifacts 4 | coverage 5 | -------------------------------------------------------------------------------- /assyst-tag/fuzz/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "assyst-tag-fuzz" 3 | version = "0.0.0" 4 | publish = false 5 | edition = "2021" 6 | 7 | [package.metadata] 8 | cargo-fuzz = true 9 | 10 | [dependencies] 11 | libfuzzer-sys = "0.4" 12 | 13 | [dependencies.assyst-tag] 14 | path = ".." 15 | 16 | # Prevent this from interfering with workspaces 17 | [workspace] 18 | members = ["."] 19 | 20 | [profile.release] 21 | debug = 1 22 | 23 | [[bin]] 24 | name = "fuzz_target_1" 25 | path = "fuzz_targets/fuzz_target_1.rs" 26 | test = false 27 | doc = false 28 | -------------------------------------------------------------------------------- /assyst-tag/fuzz/fuzz_targets/fuzz_target_1.rs: -------------------------------------------------------------------------------- 1 | #![no_main] 2 | 3 | use libfuzzer_sys::fuzz_target; 4 | 5 | fuzz_target!(|data: (&str, Vec)| { 6 | let (data, args) = data; 7 | let args = args.iter().map(|v| v.as_str()).collect::>(); 8 | if let Err(err) = assyst_tag::parse( 9 | data, 10 | &args, 11 | assyst_tag::parser::ParseMode::StopOnError, 12 | assyst_tag::NopContext, 13 | ) { 14 | assyst_tag::errors::format_error(data, err); 15 | } 16 | }); 17 | -------------------------------------------------------------------------------- /assyst-tag/src/context.rs: -------------------------------------------------------------------------------- 1 | use anyhow::anyhow; 2 | use assyst_common::eval::FakeEvalImageResponse; 3 | 4 | /// A "no-op" context, which returns an error for any of the methods 5 | /// 6 | /// This is useful for testing the parser, when you need to provide a Context but 7 | /// don't really need its functionality. 8 | pub struct NopContext; 9 | 10 | fn not_implemented() -> anyhow::Result { 11 | Err(anyhow!("Not implemented")) 12 | } 13 | 14 | /// External context for the parser 15 | /// 16 | /// It contains methods that can be provided by the caller (normally the bot crate). 17 | pub trait Context { 18 | /// Executes provided JavaScript code and returns the result (string or image) 19 | fn execute_javascript(&self, code: &str, args: Vec) -> anyhow::Result; 20 | /// Returns the URL of the last attachment 21 | fn get_last_attachment(&self) -> anyhow::Result; 22 | /// Returns the avatar URL of the provided user, or the message author 23 | fn get_avatar(&self, user_id: Option) -> anyhow::Result; 24 | /// Downloads the URL and returns the contents as a string 25 | fn download(&self, url: &str) -> anyhow::Result; 26 | /// Returns the channel ID of where this message was sent 27 | fn channel_id(&self) -> anyhow::Result; 28 | /// Returns the guild ID of where this message was sent 29 | fn guild_id(&self) -> anyhow::Result; 30 | /// Returns the user ID of the message author 31 | fn user_id(&self) -> anyhow::Result; 32 | /// Returns the tag of the provided ID 33 | fn user_tag(&self, id: Option) -> anyhow::Result; 34 | /// Loads the contents of a tag 35 | fn get_tag_contents(&self, tag: &str) -> anyhow::Result; 36 | } 37 | 38 | impl Context for NopContext { 39 | fn execute_javascript(&self, _code: &str, _args: Vec) -> anyhow::Result { 40 | not_implemented() 41 | } 42 | 43 | fn get_last_attachment(&self) -> anyhow::Result { 44 | not_implemented() 45 | } 46 | 47 | fn get_avatar(&self, _user_id: Option) -> anyhow::Result { 48 | not_implemented() 49 | } 50 | 51 | fn download(&self, _url: &str) -> anyhow::Result { 52 | not_implemented() 53 | } 54 | 55 | fn channel_id(&self) -> anyhow::Result { 56 | not_implemented() 57 | } 58 | 59 | fn guild_id(&self) -> anyhow::Result { 60 | not_implemented() 61 | } 62 | 63 | fn user_id(&self) -> anyhow::Result { 64 | not_implemented() 65 | } 66 | 67 | fn user_tag(&self, _id: Option) -> anyhow::Result { 68 | not_implemented() 69 | } 70 | 71 | fn get_tag_contents(&self, _: &str) -> anyhow::Result { 72 | not_implemented() 73 | } 74 | } 75 | 76 | impl Context for &dyn Context { 77 | fn execute_javascript(&self, code: &str, args: Vec) -> anyhow::Result { 78 | (**self).execute_javascript(code, args) 79 | } 80 | 81 | fn get_last_attachment(&self) -> anyhow::Result { 82 | (**self).get_last_attachment() 83 | } 84 | 85 | fn get_avatar(&self, user_id: Option) -> anyhow::Result { 86 | (**self).get_avatar(user_id) 87 | } 88 | 89 | fn download(&self, url: &str) -> anyhow::Result { 90 | (**self).download(url) 91 | } 92 | 93 | fn channel_id(&self) -> anyhow::Result { 94 | (**self).channel_id() 95 | } 96 | 97 | fn guild_id(&self) -> anyhow::Result { 98 | (**self).guild_id() 99 | } 100 | 101 | fn user_id(&self) -> anyhow::Result { 102 | (**self).user_id() 103 | } 104 | 105 | fn user_tag(&self, user_id: Option) -> anyhow::Result { 106 | (**self).user_tag(user_id) 107 | } 108 | 109 | fn get_tag_contents(&self, tag: &str) -> anyhow::Result { 110 | (**self).get_tag_contents(tag) 111 | } 112 | } 113 | -------------------------------------------------------------------------------- /assyst-tag/src/lib.rs: -------------------------------------------------------------------------------- 1 | #![warn(rust_2018_idioms)] 2 | #![feature(round_char_boundary, if_let_guard)] 3 | 4 | use std::cell::RefCell; 5 | use std::collections::HashMap; 6 | 7 | use assyst_common::util::filetype::Type; 8 | pub use context::{Context, NopContext}; 9 | use errors::TResult; 10 | use parser::{Counter, ParseMode, Parser, SharedState}; 11 | 12 | mod context; 13 | pub mod errors; 14 | pub mod parser; 15 | mod subtags; 16 | 17 | #[derive(Debug)] 18 | pub struct ParseResult { 19 | pub output: String, 20 | pub attachment: Option<(Vec, Type)>, 21 | } 22 | 23 | pub fn parse(input: &str, args: &[&str], mode: ParseMode, cx: C) -> TResult { 24 | let variables = RefCell::new(HashMap::new()); 25 | let counter = Counter::default(); 26 | let attachment = RefCell::new(None); 27 | let state = SharedState::new(&variables, &counter, &attachment); 28 | 29 | let output = Parser::new(input.as_bytes(), args, state, mode, &cx).parse_segment(true)?; 30 | 31 | Ok(ParseResult { 32 | output, 33 | attachment: attachment.into_inner(), 34 | }) 35 | } 36 | 37 | /// NOTE: be careful when bubbling up potential errors -- you most likely want to wrap them in 38 | /// `ErrorKind::Nested` 39 | pub fn parse_with_parent(input: &str, parent: &Parser<'_>, side_effects: bool) -> TResult { 40 | Parser::from_parent(input.as_bytes(), parent).parse_segment(side_effects) 41 | } 42 | 43 | #[cfg(test)] 44 | mod tests { 45 | use super::*; 46 | use crate::errors::ErrorKind; 47 | 48 | macro_rules! test { 49 | ($mode:expr; $( $name:ident: $input:expr => $result:pat ),+ $(,)?) => { 50 | $( 51 | #[test] 52 | fn $name() { 53 | let input = $input; 54 | let res = parse(input, &[], $mode, NopContext); 55 | assert!(matches!(res.as_ref().map_err(|err| &*err.kind).map(|ok| &*ok.output), $result)); 56 | if let Err(err) = res { 57 | 58 | // try formatting it to find any potential panic bugs 59 | errors::format_error(input, err); 60 | } 61 | } 62 | )* 63 | }; 64 | } 65 | 66 | test!(ParseMode::StopOnError; 67 | crash1: "zzzz@z{z" => Err(ErrorKind::MissingClosingBrace { .. }), 68 | crash2: "{if|z|~|\u{7}|gs|I---s{args|" => Err(ErrorKind::MissingClosingBrace{..}), 69 | crash3: "{a:a" => Err(ErrorKind::MissingClosingBrace{..}), 70 | crash4: "{note:::::JJJJ:::::::::::)[x{||z\0\0{z|||{i:||||{\u{7}Ƅ" => Err(ErrorKind::MissingClosingBrace { .. }), 71 | crash5: "{max}Ӱ______________________________" => Err(ErrorKind::ArgParseError { .. }), 72 | crash6: "{args}ӌs" => Ok("ӌs"), 73 | iter_limit: &"{max:0}".repeat(501) => Err(ErrorKind::IterLimit{..}), 74 | if_then_works: "{if:{argslen}|=|0|ok|wrong}" => Ok("ok"), 75 | if_else_works: "{if:{argslen}|=|1|wrong|ok}" => Ok("ok"), 76 | separator_outside_tag: "a|" => Ok("a|"), 77 | separator_outside_tag2: "a|}" => Ok("a|}"), 78 | separator_outside_tag3: "a|b" => Ok("a|b"), 79 | spoiler: "a||b||c" => Ok("a||b||c"), 80 | spoiler2: "a||b||" => Ok("a||b||"), 81 | spoiler_in_subparser: "{eval:a||b||c}" => Ok("a"), 82 | spoiler_in_subtag: "{note:a||b||c}" => Err(ErrorKind::MissingClosingBrace { .. }), 83 | ); 84 | 85 | test!(ParseMode::IgnoreOnError; 86 | recover1: "{foo!:42" => Ok("{foo!:42"), 87 | recover2: "{foo!:42}" => Ok("{foo!:42}"), 88 | recover3: "{foo:42}" => Ok("{foo:42}"), 89 | asynk: "(async () => {...})" => Ok("(async () => {...})"), 90 | asynk2: "(async () => { return 42 })" => Ok("(async () => { return 42 })") 91 | ); 92 | } 93 | -------------------------------------------------------------------------------- /assyst-webserver/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "assyst-webserver" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | axum-macros = { version = "0.3.0-rc.3" } 10 | axum = { version = "0.7.4", features = ["macros"] } 11 | assyst-common = { path = "../assyst-common" } 12 | assyst-database = { path = "../assyst-database" } 13 | tokio = { workspace = true } 14 | twilight-model = { workspace = true } 15 | twilight-http = { workspace = true } 16 | tracing = { workspace = true } 17 | prometheus = "0.13.3" 18 | serde = { workspace = true } 19 | lazy_static = "1.4.0" 20 | anyhow = { workspace = true } 21 | 22 | [lints] 23 | workspace = true 24 | -------------------------------------------------------------------------------- /config.template.toml: -------------------------------------------------------------------------------- 1 | ###### 2 | # NOTE: If using Patreon, you must create a `.patreon_refresh` file at the project root (i.e., same directory as this file). 3 | # This file should have you Patreon refresh token in it as a raw string. 4 | # If not using Patreon, you can safely exclude this file. However, you should also set `dev.disable_patreon_synchronisation` to true. 5 | ###### 6 | 7 | bot_id = 1234 8 | 9 | [urls] 10 | # Proxy URLs for untrusted requests. Leave blank for no proxying. 11 | proxy = [] 12 | # URL for Filer, the Assyst CDN. 13 | filer = "" 14 | # Bad translation URL. 15 | bad_translation = "" 16 | # Cobalt API instances - "primary" is optional and default false 17 | # "primary" is always the first instance that is tried, if multiple are primary it selects the lowest index primary and rest are ignored 18 | # "primary" can also be excluded for no primary instance 19 | cobalt_api = [{ url = "", key = "", primary = true }] 20 | 21 | [authentication] 22 | # Token to authenticate with Discord. 23 | discord_token = "" 24 | # Token to get paying users from Patreon. 25 | patreon_token = "" 26 | # Token to POST the stats for the bot to Top.gg 27 | top_gg_token = "" 28 | # Token that Top.gg uses to send webhooked votes to Assyst. 29 | top_gg_webhook_token = "" 30 | # Port in which the top.gg webhook runs on. 31 | top_gg_webhook_port = 3000 32 | # Authentication key for Filer, the Assyst CDN. 33 | filer_key = "" 34 | # Authentication key for NotSoAPI, for audio identification. 35 | notsoapi = "" 36 | # RapidAPI token for the `identify` command. 37 | rapidapi_token = "" 38 | 39 | # Assyst database information. 40 | [database] 41 | host = "" 42 | username = "" 43 | password = "" 44 | database = "" 45 | port = 3000 46 | 47 | [prefix] 48 | # When the bot joins a new guild, this will be the default prefix. 49 | default = "-" 50 | 51 | [logging_webhooks] 52 | panic = { token = "", id = 0 } 53 | error = { token = "", id = 0 } 54 | vote = { token = "", id = 0 } 55 | # Whether to use the webhooks on vote, panic, and error. 56 | enable_webhooks = true 57 | 58 | # Entitlements are the subscriptions for the app within Discord. You can probably leave these zeroed and the system will ignore them. 59 | [entitlements] 60 | premium_server_sku_id = 0 61 | 62 | [dev] 63 | # These Discord user IDs have full control of the bot, including developer-only commands. 64 | # Also grants max-tier premium access. 65 | admin_users = [] 66 | 67 | # When working with a dev instance, set this to the prefix "override" value for that instance 68 | # to prevent triggering the production instance and the development instance at the same time. 69 | prefix_override = "¬" 70 | 71 | # Use this for development instances to prevent the bot from attempting to process messages 72 | # in bad-translator channels. Prevents conflicts with production instance. 73 | disable_bad_translator_channels = false 74 | 75 | # Use this to disable the bot from checking reminders. Again useful when working with a 76 | # development instance to prevent conflicts. 77 | disable_reminder_check = false 78 | 79 | # Use this to top the bot from POSTing its guild and shard counts to Top.gg. 80 | # Useful for development instances. 81 | disable_bot_list_posting = false 82 | 83 | # Disables loading patrons from Patreon. 84 | disable_patreon_synchronisation = false 85 | 86 | # Whether to disable entitlement fetching. 87 | disable_entitlement_fetching = false 88 | 89 | # Whether to send the 'dev message' (see below) 90 | dev_message = false 91 | 92 | # For development instances, send a message when this guild is present in a READY event. 93 | dev_guild = 0 94 | 95 | # The channel to send the dev message in. 96 | dev_channel = 0 97 | 98 | # Override the path to the Flux executable. Useful when doing dev work on Flux. Leave blank for default. 99 | flux_executable_path_override = "" 100 | -------------------------------------------------------------------------------- /run.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash -e 2 | # script to easily start all assyst processes 3 | # not suitable for production but handy when testing 4 | 5 | trap "trap - SIGTERM && kill -- -$$" SIGINT SIGTERM EXIT 6 | 7 | mold -run cargo b -p assyst-gateway 8 | mold -run cargo b -p assyst-cache 9 | mold -run cargo b -p assyst-core 10 | 11 | cargo r -p assyst-cache & 12 | P1=$! 13 | # allow cache to start first 14 | sleep 0.5 15 | cargo r -p assyst-core & 16 | P2=$! 17 | # allow core and cache to sync before sending events 18 | sleep 0.5 19 | cargo r -p assyst-gateway & 20 | P3=$! 21 | wait $P1 $P2 $P3 -------------------------------------------------------------------------------- /rust-toolchain: -------------------------------------------------------------------------------- 1 | [toolchain] 2 | channel = "nightly-2024-09-13" 3 | components = ["cargo", "rustc", "rustc-codegen-cranelift"] -------------------------------------------------------------------------------- /rustfmt.toml: -------------------------------------------------------------------------------- 1 | max_width = 120 2 | comment_width = 100 3 | match_block_trailing_comma = true 4 | wrap_comments = true 5 | edition = "2021" 6 | error_on_line_overflow = false 7 | imports_granularity = "Module" 8 | version = "Two" 9 | ignore = [] 10 | group_imports = "StdExternalCrate" 11 | --------------------------------------------------------------------------------