Skip to main content

htb_cli/
config.rs

1use std::fs;
2use std::path::{Path, PathBuf};
3
4use serde::{Deserialize, Serialize};
5
6#[derive(Debug, Serialize, Deserialize)]
7pub struct HtbConfig {
8    #[serde(default = "default_output")]
9    pub output: String,
10
11    #[serde(default)]
12    pub vpn_server: Option<u32>,
13
14    #[serde(default)]
15    pub no_color: bool,
16
17    #[serde(default)]
18    pub cache: CacheConfig,
19}
20
21#[derive(Debug, Serialize, Deserialize)]
22pub struct CacheConfig {
23    #[serde(default = "default_cache_enabled")]
24    pub enabled: bool,
25}
26
27fn default_cache_enabled() -> bool {
28    true
29}
30
31impl Default for CacheConfig {
32    fn default() -> Self {
33        Self {
34            enabled: default_cache_enabled(),
35        }
36    }
37}
38
39fn default_output() -> String {
40    "table".into()
41}
42
43impl Default for HtbConfig {
44    fn default() -> Self {
45        Self {
46            output: default_output(),
47            vpn_server: None,
48            no_color: false,
49            cache: CacheConfig::default(),
50        }
51    }
52}
53
54impl HtbConfig {
55    pub fn load(path: Option<&Path>) -> anyhow::Result<Self> {
56        let config_path = match path {
57            Some(p) => p.to_path_buf(),
58            None => config_dir()?.join("config.toml"),
59        };
60
61        if config_path.exists() {
62            let contents = fs::read_to_string(&config_path)?;
63            Ok(toml::from_str(&contents)?)
64        } else {
65            Ok(Self::default())
66        }
67    }
68}
69
70pub fn config_dir() -> Result<PathBuf, crate::error::HtbError> {
71    dirs::home_dir()
72        .ok_or_else(|| crate::error::HtbError::Config("could not determine home directory".into()))
73        .map(|d| d.join(".htb-cli"))
74}
75
76pub fn cache_dir() -> PathBuf {
77    config_dir()
78        .map(|d| d.join("cache"))
79        .unwrap_or_else(|_| std::env::temp_dir().join("htb-cli-cache"))
80}
81
82pub fn token_path() -> Result<PathBuf, crate::error::HtbError> {
83    Ok(config_dir()?.join(".token"))
84}
85
86pub fn read_token() -> Result<String, crate::error::HtbError> {
87    let path = token_path()?;
88    if !path.exists() {
89        return Err(crate::error::HtbError::NotAuthenticated);
90    }
91    let token = fs::read_to_string(&path)?.trim().to_string();
92    if token.is_empty() {
93        return Err(crate::error::HtbError::NotAuthenticated);
94    }
95    Ok(token)
96}
97
98pub fn save_token(token: &str) -> anyhow::Result<()> {
99    let dir = config_dir()?;
100    fs::create_dir_all(&dir)?;
101
102    let path = token_path()?;
103
104    #[cfg(unix)]
105    {
106        use std::io::Write;
107        use std::os::unix::fs::OpenOptionsExt;
108        let mut f = fs::OpenOptions::new()
109            .write(true)
110            .create(true)
111            .truncate(true)
112            .mode(0o600)
113            .open(&path)?;
114        f.write_all(token.as_bytes())?;
115    }
116
117    #[cfg(not(unix))]
118    {
119        fs::write(&path, token)?;
120    }
121
122    Ok(())
123}
124
125pub fn remove_token() -> anyhow::Result<()> {
126    let path = token_path()?;
127    if path.exists() {
128        fs::remove_file(&path)?;
129    }
130    Ok(())
131}
132
133pub fn ctf_token_path() -> Result<PathBuf, crate::error::HtbError> {
134    Ok(config_dir()?.join(".ctf-token"))
135}
136
137pub fn read_ctf_token() -> Result<String, crate::error::HtbError> {
138    let path = ctf_token_path()?;
139    if !path.exists() {
140        return Err(crate::error::HtbError::NotAuthenticated);
141    }
142    let token = fs::read_to_string(&path)?.trim().to_string();
143    if token.is_empty() {
144        return Err(crate::error::HtbError::NotAuthenticated);
145    }
146    Ok(token)
147}
148
149pub fn save_ctf_token(token: &str) -> anyhow::Result<()> {
150    let dir = config_dir()?;
151    fs::create_dir_all(&dir)?;
152
153    let path = ctf_token_path()?;
154
155    #[cfg(unix)]
156    {
157        use std::io::Write;
158        use std::os::unix::fs::OpenOptionsExt;
159        let mut f = fs::OpenOptions::new()
160            .write(true)
161            .create(true)
162            .truncate(true)
163            .mode(0o600)
164            .open(&path)?;
165        f.write_all(token.as_bytes())?;
166    }
167
168    #[cfg(not(unix))]
169    {
170        fs::write(&path, token)?;
171    }
172
173    Ok(())
174}
175
176pub fn remove_ctf_token() -> anyhow::Result<()> {
177    let path = ctf_token_path()?;
178    if path.exists() {
179        fs::remove_file(&path)?;
180    }
181    Ok(())
182}
183
184#[cfg(test)]
185mod tests {
186    use super::*;
187
188    #[test]
189    fn default_config() {
190        let config = HtbConfig::default();
191        assert_eq!(config.output, "table");
192        assert_eq!(config.vpn_server, None);
193        assert!(!config.no_color);
194    }
195
196    #[test]
197    fn toml_round_trip() {
198        let config = HtbConfig {
199            output: "json".into(),
200            vpn_server: Some(1),
201            no_color: true,
202            cache: CacheConfig::default(),
203        };
204        let serialized = toml::to_string(&config).unwrap();
205        let deserialized: HtbConfig = toml::from_str(&serialized).unwrap();
206        assert_eq!(deserialized.output, "json");
207        assert_eq!(deserialized.vpn_server, Some(1));
208        assert!(deserialized.no_color);
209    }
210
211    #[test]
212    fn partial_toml_uses_defaults() {
213        let input = r#"output = "json""#;
214        let config: HtbConfig = toml::from_str(input).unwrap();
215        assert_eq!(config.output, "json");
216        assert_eq!(config.vpn_server, None);
217        assert!(!config.no_color);
218    }
219
220    #[test]
221    fn empty_toml_uses_defaults() {
222        let config: HtbConfig = toml::from_str("").unwrap();
223        assert_eq!(config.output, "table");
224    }
225
226    #[test]
227    fn token_round_trip() {
228        let dir = std::env::temp_dir().join("htb-cli-test-token");
229        let _ = std::fs::create_dir_all(&dir);
230        let path = dir.join(".token");
231
232        std::fs::write(&path, "test-token-123").unwrap();
233        let read_back = std::fs::read_to_string(&path).unwrap();
234        assert_eq!(read_back.trim(), "test-token-123");
235
236        let _ = std::fs::remove_dir_all(&dir);
237    }
238
239    #[test]
240    fn ctf_token_path_ends_with_ctf_token() {
241        let path = ctf_token_path().unwrap();
242        assert!(path.ends_with(".ctf-token"));
243    }
244
245    #[test]
246    fn ctf_token_path_differs_from_labs_token() {
247        let labs = token_path().unwrap();
248        let ctf = ctf_token_path().unwrap();
249        assert_ne!(labs, ctf);
250    }
251}