Skip to main content

kovi_onebot/driver/
config.rs

1use dialoguer::theme::ColorfulTheme;
2use dialoguer::{Input, Select};
3use kovi::error::BotBuildError;
4use serde::{Deserialize, Serialize};
5use std::fmt::Display;
6use std::fs;
7use std::io::Write as _;
8use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
9use std::path::Path;
10
11#[derive(Debug, Clone, Deserialize, Serialize)]
12pub struct OneBotDriverConfig {
13    pub server: Server,
14}
15
16impl OneBotDriverConfig {
17    pub fn normalize_path(self) -> Self {
18        Self {
19            server: Server {
20                host: self.server.host,
21                port: self.server.port,
22                access_token: self.server.access_token,
23                secure: self.server.secure,
24                path: if self.server.path.ends_with('/') {
25                    self.server.path
26                } else {
27                    format!("{}/", self.server.path)
28                },
29                all_in_one: self.server.all_in_one,
30            },
31        }
32    }
33}
34
35/// server信息
36#[derive(Deserialize, Serialize, Debug, Clone)]
37pub struct Server {
38    pub host: Host,
39    pub port: u16,
40    pub access_token: String,
41    pub secure: bool,
42    /// path route to ws
43    #[serde(default = "default_path")]
44    pub path: String,
45
46    /// all in one single "/" endpoint
47    #[serde(default)]
48    pub all_in_one: bool,
49}
50
51/// when not specified, use "/" instead.
52fn default_path() -> String {
53    "/".into()
54}
55
56impl Server {
57    pub fn new(
58        host: Host,
59        port: u16,
60        access_token: String,
61        secure: bool,
62        path: String,
63        all_in_one: bool,
64    ) -> Self {
65        Server {
66            host,
67            port,
68            access_token,
69            secure,
70            path,
71            all_in_one,
72        }
73    }
74}
75
76impl Server {
77    /// 根据 path 后缀构建 WebSocket URL,例如 `ws_url("api")` → `ws://host:port/api`
78    /// 如果启用了 all_in_one 模式,path 将被忽略
79    pub fn ws_url(&self, path: &str) -> String {
80        let path = if self.all_in_one {
81            "".to_string()
82        } else {
83            format!("/{}", path)
84        };
85
86        let protocol = if self.secure { "wss" } else { "ws" };
87        let host = match &self.host {
88            Host::IpAddr(std::net::IpAddr::V6(ip)) => format!("[{ip}]"),
89            Host::IpAddr(ip) => ip.to_string(),
90            Host::Domain(d) => d.clone(),
91        };
92
93        format!(
94            "{protocol}://{host}:{self_port}{self_path}{path}",
95            self_port = self.port,
96            self_path = match self.path.as_str() {
97                "" => String::new(),
98                p => p.to_string(),
99            },
100        )
101    }
102}
103
104impl AsRef<OneBotDriverConfig> for OneBotDriverConfig {
105    fn as_ref(&self) -> &OneBotDriverConfig {
106        self
107    }
108}
109
110#[derive(Deserialize, Serialize, Debug, Clone)]
111#[serde(untagged)]
112pub enum Host {
113    IpAddr(IpAddr),
114    Domain(String),
115}
116
117impl Display for Host {
118    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
119        match self {
120            Host::IpAddr(ip) => write!(f, "{ip}"),
121            Host::Domain(domain) => write!(f, "{domain}"),
122        }
123    }
124}
125
126/// 将配置文件写入磁盘
127fn config_file_write_and_return(file_path: &Path) -> Result<OneBotDriverConfig, std::io::Error> {
128    enum HostType {
129        IPv4,
130        IPv6,
131        Domain,
132    }
133
134    let host_type: HostType = {
135        let items = ["IPv4", "IPv6", "Domain"];
136        let select = Select::with_theme(&ColorfulTheme::default())
137            .with_prompt("What is the type of the host of the OneBot server?")
138            .items(&items)
139            .default(0)
140            .interact()
141            .expect("unreachable");
142
143        match select {
144            0 => HostType::IPv4,
145            1 => HostType::IPv6,
146            2 => HostType::Domain,
147            _ => panic!(), // 不可能的事情
148        }
149    };
150
151    let host = match host_type {
152        HostType::IPv4 => {
153            let ip = Input::with_theme(&ColorfulTheme::default())
154                .with_prompt("What is the IP of the OneBot server?")
155                .default(Ipv4Addr::new(127, 0, 0, 1))
156                .interact_text()
157                .expect("unreachable");
158            Host::IpAddr(IpAddr::V4(ip))
159        }
160        HostType::IPv6 => {
161            let ip = Input::with_theme(&ColorfulTheme::default())
162                .with_prompt("What is the IP of the OneBot server?")
163                .default(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1))
164                .interact_text()
165                .expect("unreachable");
166            Host::IpAddr(IpAddr::V6(ip))
167        }
168        HostType::Domain => {
169            let domain = Input::with_theme(&ColorfulTheme::default())
170                .with_prompt("What is the domain of the OneBot server?")
171                .default("localhost".to_string())
172                .interact_text()
173                .expect("unreachable");
174            Host::Domain(domain)
175        }
176    };
177
178    let port: u16 = Input::with_theme(&ColorfulTheme::default())
179        .with_prompt("What is the port of the OneBot server?")
180        .default(8081)
181        .interact_text()
182        .expect("unreachable");
183
184    let access_token: String = Input::with_theme(&ColorfulTheme::default())
185        .with_prompt("What is the access_token of the OneBot server? (Optional)")
186        .default("".to_string())
187        .show_default(false)
188        .interact_text()
189        .expect("unreachable");
190
191    let path: String = Input::with_theme(&ColorfulTheme::default())
192        .with_prompt("What is the route path of websocket server?")
193        .default("/".to_string())
194        .interact_text()
195        .expect("unreachable");
196
197    let more: bool = {
198        let items = ["No", "Yes"];
199        let select = Select::with_theme(&ColorfulTheme::default())
200            .with_prompt("Do you want to view more optional options?")
201            .items(&items)
202            .default(0)
203            .interact()
204            .expect("unreachable");
205
206        match select {
207            0 => false,
208            1 => true,
209            _ => unreachable!(),
210        }
211    };
212
213    let mut secure = false;
214    let mut all_in_one = false;
215    if more {
216        fn select_bool(prompt: &str) -> bool {
217            let items = ["No", "Yes"];
218            let select = Select::with_theme(&ColorfulTheme::default())
219                .with_prompt(prompt)
220                .items(&items)
221                .default(0)
222                .interact()
223                .expect("unreachable");
224
225            select == 1
226        }
227        secure = select_bool("Enable secure connection? (WSS)");
228        all_in_one = select_bool("Use single ws api endpoint?");
229    }
230
231    let config = OneBotDriverConfig {
232        server: Server {
233            host,
234            port,
235            access_token,
236            secure,
237            path,
238            all_in_one,
239        },
240    };
241
242    let mut doc = match fs::read_to_string(file_path) {
243        Ok(content) => match content.parse::<toml_edit::DocumentMut>() {
244            Ok(d) => d,
245            Err(err) => {
246                eprintln!(
247                    "Failed to parse existing config, creating new document: {}",
248                    err
249                );
250                toml_edit::DocumentMut::new()
251            }
252        },
253        Err(_) => toml_edit::DocumentMut::new(),
254    };
255
256    doc["server"] = toml_edit::table();
257    doc["server"]["host"] = match &config.server.host {
258        Host::IpAddr(ip) => toml_edit::value(ip.to_string()),
259        Host::Domain(domain) => toml_edit::value(domain),
260    };
261    doc["server"]["port"] = toml_edit::value(config.server.port as i64);
262    doc["server"]["access_token"] = toml_edit::value(&config.server.access_token);
263    doc["server"]["secure"] = toml_edit::value(config.server.secure);
264    doc["server"]["path"] = toml_edit::value(&config.server.path);
265    doc["server"]["all_in_one"] = toml_edit::value(config.server.all_in_one);
266
267    let file = fs::File::create(file_path)?;
268    let mut writer = std::io::BufWriter::new(file);
269    writer.write_all(doc.to_string().as_bytes())?;
270
271    Ok(config)
272}
273
274/// 读取本地Kovi.conf.toml文件
275pub fn load_local_conf() -> Result<OneBotDriverConfig, BotBuildError> {
276    let path = Path::new("kovi.conf.toml");
277    let kovi_conf_file_exist = fs::metadata(path).is_ok();
278
279    #[derive(Deserialize, Serialize, Debug, Clone)]
280    struct TempKoviConfig {
281        server: Option<Server>,
282    }
283
284    let conf_json: OneBotDriverConfig = if kovi_conf_file_exist {
285        match fs::read_to_string(path) {
286            Ok(v) => match toml::from_str::<TempKoviConfig>(&v) {
287                Ok(conf) => match conf.server {
288                    Some(server) => OneBotDriverConfig { server },
289                    None => config_file_write_and_return(path)
290                        .map_err(|e| BotBuildError::FileCreateError(e.to_string()))?,
291                },
292                Err(err) => {
293                    eprintln!("Configuration file parsing error: {err}");
294                    config_file_write_and_return(path)
295                        .map_err(|e| BotBuildError::FileCreateError(e.to_string()))?
296                }
297            },
298            Err(err) => {
299                return Err(BotBuildError::FileReadError(err.to_string()));
300            }
301        }
302    } else {
303        config_file_write_and_return(path)
304            .map_err(|e| BotBuildError::FileCreateError(e.to_string()))?
305    };
306
307    Ok(conf_json)
308}