kovi/
bot.rs

1use ahash::{HashMapExt as _, RandomState};
2use dialoguer::theme::ColorfulTheme;
3use dialoguer::{Input, Select};
4use plugin_builder::Listen;
5use serde::{Deserialize, Serialize};
6use serde_json::{self, Value};
7use std::collections::{HashMap, HashSet};
8use std::env;
9use std::fmt::{Debug, Display};
10use std::future::Future;
11use std::io::Write as _;
12use std::net::{Ipv4Addr, Ipv6Addr};
13use std::pin::Pin;
14use std::{fs, net::IpAddr, sync::Arc};
15use tokio::sync::mpsc::{self};
16use tokio::sync::{oneshot, watch};
17use tokio::task::JoinHandle;
18
19use crate::error::{BotBuildError, BotError};
20use crate::task::TASK_MANAGER;
21
22#[cfg(feature = "plugin-access-control")]
23pub use crate::bot::runtimebot::kovi_api::AccessControlMode;
24
25pub(crate) mod connect;
26pub(crate) mod handler;
27pub(crate) mod run;
28
29pub mod message;
30pub mod plugin_builder;
31pub mod runtimebot;
32
33tokio::task_local! {
34    pub static PLUGIN_BUILDER: crate::PluginBuilder;
35}
36
37tokio::task_local! {
38    pub static PLUGIN_NAME: Arc<String>;
39}
40
41/// kovi的配置
42#[derive(Debug, Clone, Deserialize, Serialize)]
43pub struct KoviConf {
44    pub config: Config,
45    pub server: Server,
46}
47
48impl AsRef<KoviConf> for KoviConf {
49    fn as_ref(&self) -> &KoviConf {
50        self
51    }
52}
53
54#[derive(Debug, Clone, Deserialize, Serialize)]
55pub struct Config {
56    pub main_admin: i64,
57    pub admins: Vec<i64>,
58    pub debug: bool,
59}
60
61impl KoviConf {
62    pub fn new(main_admin: i64, admins: Option<Vec<i64>>, server: Server, debug: bool) -> Self {
63        KoviConf {
64            config: Config {
65                main_admin,
66                admins: admins.unwrap_or_default(),
67                debug,
68            },
69            server,
70        }
71    }
72}
73
74type KoviAsyncFn = dyn Fn() -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync;
75
76impl Drop for Bot {
77    fn drop(&mut self) {
78        for i in self.run_abort.iter() {
79            i.abort();
80        }
81    }
82}
83
84/// bot结构体
85#[derive(Clone)]
86pub struct Bot {
87    pub information: BotInformation,
88    pub(crate) plugins: HashMap<String, BotPlugin, RandomState>,
89    pub(crate) run_abort: Vec<tokio::task::AbortHandle>,
90}
91
92#[derive(Clone)]
93pub(crate) struct BotPlugin {
94    pub(crate) enable_on_startup: bool,
95    pub(crate) enabled: watch::Sender<bool>,
96
97    pub(crate) name: String,
98    pub(crate) version: String,
99    pub(crate) main: Arc<KoviAsyncFn>,
100    pub(crate) listen: Listen,
101
102    #[cfg(feature = "plugin-access-control")]
103    pub(crate) access_control: bool,
104    #[cfg(feature = "plugin-access-control")]
105    pub(crate) list_mode: AccessControlMode,
106    #[cfg(feature = "plugin-access-control")]
107    pub(crate) access_list: AccessList,
108}
109
110#[cfg(feature = "plugin-access-control")]
111#[derive(Clone, Debug, Default, Deserialize, Serialize)]
112pub struct AccessList {
113    pub friends: HashSet<i64>,
114    pub groups: HashSet<i64>,
115}
116
117#[derive(Clone, Debug, Deserialize, Serialize)]
118pub struct PluginInfo {
119    pub name: String,
120    pub version: String,
121    /// 插件是否启用
122    pub enabled: bool,
123    /// 插件是否在Bot启动时启用
124    pub enable_on_startup: bool,
125    /// 插件是否启用框架级访问控制
126    #[cfg(feature = "plugin-access-control")]
127    pub access_control: bool,
128    /// 插件的访问控制模式
129    #[cfg(feature = "plugin-access-control")]
130    pub list_mode: AccessControlMode,
131    /// 插件的访问控制列表
132    #[cfg(feature = "plugin-access-control")]
133    pub access_list: AccessList,
134}
135
136#[derive(Debug, Clone, Deserialize, Serialize)]
137struct PluginStatus {
138    enable_on_startup: bool,
139    #[cfg(feature = "plugin-access-control")]
140    access_control: bool,
141    #[cfg(feature = "plugin-access-control")]
142    list_mode: AccessControlMode,
143    #[cfg(feature = "plugin-access-control")]
144    access_list: AccessList,
145}
146
147/// bot信息结构体
148#[derive(Debug, Clone)]
149pub struct BotInformation {
150    pub main_admin: i64,
151    pub deputy_admins: HashSet<i64>,
152    pub server: Server,
153}
154/// server信息
155#[derive(Deserialize, Serialize, Debug, Clone)]
156pub struct Server {
157    pub host: Host,
158    pub port: u16,
159    pub access_token: String,
160    pub secure: bool,
161}
162
163#[derive(Deserialize, Serialize, Debug, Clone)]
164#[serde(untagged)]
165pub enum Host {
166    IpAddr(IpAddr),
167    Domain(String),
168}
169
170impl Display for Host {
171    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
172        match self {
173            Host::IpAddr(ip) => write!(f, "{}", ip),
174            Host::Domain(domain) => write!(f, "{}", domain),
175        }
176    }
177}
178
179impl Server {
180    pub fn new(host: Host, port: u16, access_token: String, secure: bool) -> Self {
181        Server {
182            host,
183            port,
184            access_token,
185            secure,
186        }
187    }
188}
189
190#[derive(Debug, Deserialize, Serialize, Clone)]
191pub struct SendApi {
192    pub action: String,
193    pub params: Value,
194    pub echo: String,
195}
196
197#[derive(Debug, Clone, Deserialize, Serialize)]
198pub struct ApiReturn {
199    pub status: String,
200    pub retcode: i32,
201    pub data: Value,
202    pub echo: String,
203}
204
205pub(crate) type ApiAndOneshot = (
206    SendApi,
207    Option<oneshot::Sender<Result<ApiReturn, ApiReturn>>>,
208);
209
210impl std::fmt::Display for ApiReturn {
211    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
212        write!(
213            f,
214            "status: {}, retcode: {}, data: {}, echo: {}",
215            self.status, self.retcode, self.data, self.echo
216        )
217    }
218}
219
220impl std::fmt::Display for SendApi {
221    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
222        write!(f, "{}", serde_json::to_string(self).unwrap())
223    }
224}
225
226impl SendApi {
227    pub fn new(action: &str, params: Value, echo: &str) -> Self {
228        SendApi {
229            action: action.to_string(),
230            params,
231            echo: echo.to_string(),
232        }
233    }
234}
235
236impl BotPlugin {
237    fn shutdown(&mut self) -> JoinHandle<()> {
238        log::debug!("Plugin '{}' is dropping.", self.name,);
239
240        let plugin_name_ = Arc::new(self.name.clone());
241
242        let mut task_vec = Vec::new();
243
244        for listen in &self.listen.drop {
245            let listen_clone = listen.clone();
246            let plugin_name_ = plugin_name_.clone();
247            let task = tokio::spawn(async move {
248                PLUGIN_NAME
249                    .scope(plugin_name_, Bot::handler_drop(listen_clone))
250                    .await;
251            });
252            task_vec.push(task);
253        }
254
255        TASK_MANAGER.disable_plugin(&self.name);
256
257        self.enabled.send_modify(|v| {
258            *v = false;
259        });
260        self.listen.clear();
261        tokio::spawn(async move {
262            for task in task_vec {
263                let _ = task.await;
264            }
265        })
266    }
267}
268
269impl Bot {
270    /// 构建一个bot实例
271    /// # Examples
272    /// ```
273    /// let conf = KoviConf::new(
274    ///     123456,
275    ///     None,
276    ///     Server {
277    ///         host: "127.0.0.1".parse(),
278    ///         port: 8081,
279    ///         access_token: "",
280    ///     },
281    ///     false,
282    ///     None,
283    /// );
284    /// let bot = Bot::build(conf);
285    /// bot.run()
286    /// ```
287    pub fn build<C>(conf: C) -> Bot
288    where
289        C: AsRef<KoviConf>,
290    {
291        let conf = conf.as_ref();
292        Bot {
293            information: BotInformation {
294                main_admin: conf.config.main_admin,
295                deputy_admins: conf.config.admins.iter().cloned().collect(),
296                server: conf.server.clone(),
297            },
298            plugins: HashMap::<_, _, RandomState>::new(),
299            run_abort: Vec::new(),
300        }
301    }
302
303    /// 挂载插件的启动函数。
304    pub fn mount_main<T>(&mut self, name: T, version: T, main: Arc<KoviAsyncFn>)
305    where
306        String: From<T>,
307    {
308        let name = String::from(name);
309        let version = String::from(version);
310        let (tx, _rx) = watch::channel(true);
311        let bot_plugin = BotPlugin {
312            enable_on_startup: true,
313            enabled: tx,
314            name: name.clone(),
315            version,
316            main,
317            listen: Listen::default(),
318
319            #[cfg(feature = "plugin-access-control")]
320            access_control: false,
321            #[cfg(feature = "plugin-access-control")]
322            list_mode: AccessControlMode::WhiteList,
323            #[cfg(feature = "plugin-access-control")]
324            access_list: AccessList::default(),
325        };
326        self.plugins.insert(name, bot_plugin);
327    }
328
329    /// 读取本地Kovi.conf.toml文件
330    pub fn load_local_conf() -> Result<KoviConf, BotBuildError> {
331        //检测文件是kovi.conf.json还是kovi.conf.toml
332        let kovi_conf_file_exist = fs::metadata("kovi.conf.toml").is_ok();
333
334        let conf_json: KoviConf = if kovi_conf_file_exist {
335            match fs::read_to_string("kovi.conf.toml") {
336                Ok(v) => match toml::from_str(&v) {
337                    Ok(conf) => conf,
338                    Err(err) => {
339                        eprintln!("Configuration file parsing error: {}", err);
340                        config_file_write_and_return()
341                            .map_err(|e| BotBuildError::FileCreateError(e.to_string()))?
342                    }
343                },
344                Err(err) => {
345                    return Err(BotBuildError::FileReadError(err.to_string()));
346                }
347            }
348        } else {
349            config_file_write_and_return()
350                .map_err(|e| BotBuildError::FileCreateError(e.to_string()))?
351        };
352
353        unsafe {
354            if env::var("RUST_LOG").is_err() {
355                if conf_json.config.debug {
356                    env::set_var("RUST_LOG", "debug");
357                } else {
358                    env::set_var("RUST_LOG", "info");
359                }
360            }
361        }
362
363        Ok(conf_json)
364    }
365}
366
367impl Bot {
368    /// 使用KoviConf设置插件在Bot启动时的状态
369    ///
370    /// 如果配置文件中没有对应的插件，将会被忽略，保留插件默认状态
371    ///
372    /// 如果配置文件读取失败或者解析toml失败，将会保留插件默认状态
373    pub fn set_plugin_startup_use_file(mut self) -> Self {
374        let file_path = "kovi.plugin.toml";
375        let content = match fs::read_to_string(file_path) {
376            Ok(v) => {
377                log::debug!("Set plugin startup use file successfully");
378                v
379            }
380            Err(e) => {
381                log::debug!("Failed to read file: {}", e);
382                return self;
383            }
384        };
385        let mut plugin_status_map: HashMap<String, PluginStatus> = match toml::from_str(&content) {
386            Ok(v) => v,
387            Err(e) => {
388                log::debug!("Failed to parse toml: {}", e);
389                return self;
390            }
391        };
392
393        for (name, plugin) in self.plugins.iter_mut() {
394            if let Some(plugin_status) = plugin_status_map.remove(name) {
395                plugin.enable_on_startup = plugin_status.enable_on_startup;
396                plugin.enabled.send_modify(|v| {
397                    *v = plugin_status.enable_on_startup;
398                });
399                #[cfg(feature = "plugin-access-control")]
400                {
401                    plugin.access_control = plugin_status.access_control;
402                    plugin.list_mode = plugin_status.list_mode;
403                    plugin.access_list = plugin_status.access_list;
404                }
405            }
406        }
407
408        self
409    }
410
411    /// 使用KoviConf设置插件在Bot启动时的状态
412    ///
413    /// 如果配置文件中没有对应的插件，将会被忽略，保留插件默认状态
414    ///
415    /// 如果配置文件读取失败或者解析toml失败，将会保留插件默认状态
416    pub fn set_plugin_startup_use_file_ref(&mut self) {
417        let file_path = "kovi.plugin.toml";
418        let content = match fs::read_to_string(file_path) {
419            Ok(v) => {
420                log::debug!("Set plugin startup use file successfully");
421                v
422            }
423            Err(e) => {
424                log::debug!("Failed to read file: {}", e);
425                return;
426            }
427        };
428        let mut plugin_status_map: HashMap<String, PluginStatus> = match toml::from_str(&content) {
429            Ok(v) => v,
430            Err(e) => {
431                log::debug!("Failed to parse toml: {}", e);
432                return;
433            }
434        };
435
436        for (name, plugin) in self.plugins.iter_mut() {
437            if let Some(plugin_status) = plugin_status_map.remove(name) {
438                plugin.enable_on_startup = plugin_status.enable_on_startup;
439                plugin.enabled.send_modify(|v| {
440                    *v = plugin_status.enable_on_startup;
441                });
442                #[cfg(feature = "plugin-access-control")]
443                {
444                    plugin.access_control = plugin_status.access_control;
445                    plugin.list_mode = plugin_status.list_mode;
446                    plugin.access_list = plugin_status.access_list;
447                }
448            }
449        }
450    }
451
452    /// 设置全部插件在Bot启动时的状态
453    pub fn set_all_plugin_startup(mut self, enabled: bool) -> Self {
454        for plugin in self.plugins.values_mut() {
455            plugin.enable_on_startup = enabled;
456            plugin.enabled.send_modify(|v| {
457                *v = enabled;
458            });
459        }
460        self
461    }
462
463    /// 设置全部插件在Bot启动时的状态
464    pub fn set_all_plugin_startup_ref(&mut self, enabled: bool) {
465        for plugin in self.plugins.values_mut() {
466            plugin.enable_on_startup = enabled;
467            plugin.enabled.send_modify(|v| {
468                *v = enabled;
469            });
470        }
471    }
472
473    /// 设置单个插件在Bot启动时的状态
474    pub fn set_plugin_startup<T: AsRef<str>>(
475        mut self,
476        name: T,
477        enabled: bool,
478    ) -> Result<Self, BotError> {
479        let name = name.as_ref();
480        if let Some(plugin) = self.plugins.get_mut(name) {
481            plugin.enable_on_startup = enabled;
482            plugin.enabled.send_modify(|v| {
483                *v = enabled;
484            });
485            Ok(self)
486        } else {
487            Err(BotError::PluginNotFound(format!(
488                "Plugin {} not found",
489                name
490            )))
491        }
492    }
493
494    /// 设置单个插件在Bot启动时的状态
495    pub fn set_plugin_startup_ref<T: AsRef<str>>(
496        &mut self,
497        name: T,
498        enabled: bool,
499    ) -> Result<(), BotError> {
500        let name = name.as_ref();
501        if let Some(plugin) = self.plugins.get_mut(name) {
502            plugin.enable_on_startup = enabled;
503            plugin.enabled.send_modify(|v| {
504                *v = enabled;
505            });
506            Ok(())
507        } else {
508            Err(BotError::PluginNotFound(format!(
509                "Plugin {} not found",
510                name
511            )))
512        }
513    }
514
515    #[cfg(any(feature = "save_plugin_status", feature = "save_bot_admin"))]
516    pub(crate) fn save_bot_status(&self) {
517        #[cfg(feature = "save_plugin_status")]
518        {
519            let _file_path = "kovi.plugin.toml";
520
521            let mut plugin_status = HashMap::new();
522            for (name, plugin) in self.plugins.iter() {
523                plugin_status.insert(name.clone(), PluginStatus {
524                    enable_on_startup: *plugin.enabled.borrow(),
525                    #[cfg(feature = "plugin-access-control")]
526                    access_control: plugin.access_control,
527                    #[cfg(feature = "plugin-access-control")]
528                    list_mode: plugin.list_mode,
529                    #[cfg(feature = "plugin-access-control")]
530                    access_list: plugin.access_list.clone(),
531                });
532            }
533
534            let serialized = match toml::to_string(&plugin_status) {
535                Ok(s) => s,
536                Err(e) => {
537                    log::error!("Failed to serialize plugin status: {}", e);
538                    return;
539                }
540            };
541            if let Err(e) = fs::write(_file_path, serialized) {
542                log::error!("Failed to write plugin status to file: {}", e);
543            }
544        }
545
546        #[cfg(feature = "save_bot_admin")]
547        {
548            let file_path = "kovi.conf.toml";
549            let existing_content = fs::read_to_string(file_path).unwrap_or_default();
550
551            let mut doc = existing_content
552                .parse::<toml_edit::DocumentMut>()
553                .unwrap_or_else(|_| toml_edit::DocumentMut::new());
554
555            // 确保 "config" 存在
556            if !doc.contains_key("config") {
557                doc["config"] = toml_edit::table();
558            }
559
560            // 更新 "config" 中的 admin 信息
561            doc["config"]["main_admin"] = toml_edit::value(self.information.main_admin);
562            doc["config"]["admins"] = toml_edit::Item::Value(toml_edit::Value::Array(
563                self.information
564                    .deputy_admins
565                    .iter()
566                    .map(|&x| toml_edit::Value::from(x))
567                    .collect(),
568            ));
569
570            match fs::File::create(file_path) {
571                Ok(file) => {
572                    let mut writer = std::io::BufWriter::new(file);
573                    if let Err(e) = writer.write_all(doc.to_string().as_bytes()) {
574                        log::error!("Failed to write to file: {}", e);
575                    }
576                }
577                Err(e) => {
578                    log::error!("Failed to create file: {}", e);
579                }
580            }
581        }
582    }
583}
584
585/// 将配置文件写入磁盘
586fn config_file_write_and_return() -> Result<KoviConf, std::io::Error> {
587    enum HostType {
588        IPv4,
589        IPv6,
590        Domain,
591    }
592
593    let host_type: HostType = {
594        let items = ["IPv4", "IPv6", "Domain"];
595        let select = Select::with_theme(&ColorfulTheme::default())
596            .with_prompt("What is the type of the host of the OneBot server?")
597            .items(&items)
598            .default(0)
599            .interact()
600            .unwrap();
601
602        match select {
603            0 => HostType::IPv4,
604            1 => HostType::IPv6,
605            2 => HostType::Domain,
606            _ => panic!(), //不可能的事情
607        }
608    };
609
610    let host = match host_type {
611        HostType::IPv4 => {
612            let ip = Input::with_theme(&ColorfulTheme::default())
613                .with_prompt("What is the IP of the OneBot server?")
614                .default(Ipv4Addr::new(127, 0, 0, 1))
615                .interact_text()
616                .unwrap();
617            Host::IpAddr(IpAddr::V4(ip))
618        }
619        HostType::IPv6 => {
620            let ip = Input::with_theme(&ColorfulTheme::default())
621                .with_prompt("What is the IP of the OneBot server?")
622                .default(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1))
623                .interact_text()
624                .unwrap();
625            Host::IpAddr(IpAddr::V6(ip))
626        }
627        HostType::Domain => {
628            let domain = Input::with_theme(&ColorfulTheme::default())
629                .with_prompt("What is the domain of the OneBot server?")
630                .default("localhost".to_string())
631                .interact_text()
632                .unwrap();
633            Host::Domain(domain)
634        }
635    };
636
637    let port: u16 = Input::with_theme(&ColorfulTheme::default())
638        .with_prompt("What is the port of the OneBot server?")
639        .default(8081)
640        .interact_text()
641        .unwrap();
642
643    let access_token: String = Input::with_theme(&ColorfulTheme::default())
644        .with_prompt("What is the access_token of the OneBot server? (Optional)")
645        .default("".to_string())
646        .show_default(false)
647        .interact_text()
648        .unwrap();
649
650    let main_admin: i64 = Input::with_theme(&ColorfulTheme::default())
651        .with_prompt("What is the ID of the main administrator? (Not used yet)")
652        .allow_empty(true)
653        .interact_text()
654        .unwrap();
655
656    // 是否查看更多可选选项
657    let more: bool = {
658        let items = ["No", "Yes"];
659        let select = Select::with_theme(&ColorfulTheme::default())
660            .with_prompt("Do you want to view more optional options?")
661            .items(&items)
662            .default(0)
663            .interact()
664            .unwrap();
665
666        match select {
667            0 => false,
668            1 => true,
669            _ => panic!(), //不可能的事情
670        }
671    };
672
673    let mut secure = false;
674    if more {
675        // wss https? tls?
676        secure = {
677            let items = vec!["No", "Yes"];
678            let select = Select::with_theme(&ColorfulTheme::default())
679                // .with_prompt("Enable secure connection? (HTTPS/WSS)")
680                .with_prompt("Enable secure connection? (WSS)")
681                .items(&items)
682                .default(0)
683                .interact()
684                .unwrap();
685
686            match select {
687                0 => false,
688                1 => true,
689                _ => panic!(), //不可能的事情
690            }
691        };
692    }
693
694    let config = KoviConf::new(
695        main_admin,
696        None,
697        Server::new(host, port, access_token, secure),
698        false,
699    );
700
701    let mut doc = toml_edit::DocumentMut::new();
702    doc["config"] = toml_edit::table();
703    doc["config"]["main_admin"] = toml_edit::value(config.config.main_admin);
704    doc["config"]["admins"] = toml_edit::Item::Value(toml_edit::Value::Array(
705        config
706            .config
707            .admins
708            .iter()
709            .map(|&x| toml_edit::Value::from(x))
710            .collect(),
711    ));
712    doc["config"]["debug"] = toml_edit::value(config.config.debug);
713
714    doc["server"] = toml_edit::table();
715    doc["server"]["host"] = match &config.server.host {
716        Host::IpAddr(ip) => toml_edit::value(ip.to_string()),
717        Host::Domain(domain) => toml_edit::value(domain.clone()),
718    };
719    doc["server"]["port"] = toml_edit::value(config.server.port as i64);
720    doc["server"]["access_token"] = toml_edit::value(config.server.access_token.clone());
721    doc["server"]["secure"] = toml_edit::value(config.server.secure);
722
723    let file = fs::File::create("kovi.conf.toml")?;
724    let mut writer = std::io::BufWriter::new(file);
725    writer.write_all(doc.to_string().as_bytes())?;
726
727    Ok(config)
728}
729
730#[macro_export]
731macro_rules! build_bot {
732    ($( $plugin:ident ),* $(,)* ) => {
733        {
734            let conf = match kovi::bot::Bot::load_local_conf() {
735                Ok(c) => c,
736                Err(e) => {
737                    eprintln!("Error loading config: {}", e);
738                    panic!("Failed to load config");
739                }
740            };
741            kovi::logger::try_set_logger();
742            let mut bot = kovi::bot::Bot::build(&conf);
743
744            $(
745                let (crate_name, crate_version) = $plugin::__kovi_get_plugin_info();
746                kovi::log::info!("Mounting plugin: {}", crate_name);
747                bot.mount_main(crate_name, crate_version, std::sync::Arc::new($plugin::__kovi_run_async_plugin));
748            )*
749
750            bot.set_plugin_startup_use_file_ref();
751            bot
752        }
753    };
754}
755
756#[test]
757fn build_bot() {
758    let conf = KoviConf::new(
759        123456,
760        None,
761        Server {
762            host: Host::IpAddr("127.0.0.1".parse().unwrap()),
763            port: 8081,
764            access_token: "".to_string(),
765            secure: false,
766        },
767        false,
768    );
769    let _ = Bot::build(conf);
770}
kovi/bot.rs

kovi/
bot.rs
