1use std::collections::BTreeMap;
2use std::process::{Command, Stdio};
3
4use crate::error::Result;
5use crate::models::{
6 MSFHostCheckResult, MSFHostInfo, MSFHostsResult, MSFLoginResult, MSFWorkspaceInfo,
7 MSFWorkspacesResult,
8};
9use crate::msgpack::{decode, encode, MsgValue};
10
11#[derive(Debug, Clone)]
12pub struct MsfRpcClient {
13 host: String,
14 port: u16,
15 ssl: bool,
16 username: String,
17 password: String,
18 token: Option<String>,
19}
20
21impl MsfRpcClient {
22 pub fn new(host: &str, port: u16, username: &str, password: &str, ssl: bool) -> Self {
23 Self {
24 host: host.to_string(),
25 port,
26 ssl,
27 username: username.to_string(),
28 password: password.to_string(),
29 token: None,
30 }
31 }
32
33 pub fn login(&mut self) -> Result<MSFLoginResult> {
34 let url = build_url(&self.host, self.port, self.ssl);
35 let payload = encode(&msg_map([
36 ("method", msg_str("auth.login")),
37 (
38 "params",
39 msg_array(vec![msg_str(&self.username), msg_str(&self.password)]),
40 ),
41 ]));
42
43 match send_request(&url, &payload, 10) {
44 Ok(val) => {
45 let success = val.get("result").and_then(MsgValue::as_str) == Some("success");
46 let token = val
47 .get("token")
48 .and_then(MsgValue::as_str)
49 .unwrap_or("")
50 .to_string();
51
52 if success && !token.is_empty() {
53 self.token = Some(token.clone());
54 Ok(MSFLoginResult {
55 success: true,
56 token,
57 error: String::new(),
58 })
59 } else {
60 Ok(MSFLoginResult {
61 success: false,
62 token: String::new(),
63 error: "Authentication failed".to_string(),
64 })
65 }
66 }
67 Err(error) => Ok(MSFLoginResult {
68 success: false,
69 token: String::new(),
70 error,
71 }),
72 }
73 }
74
75 pub fn logout(&mut self) {
76 if let Some(token) = self.token.clone() {
77 let url = build_url(&self.host, self.port, self.ssl);
78 let payload = encode(&msg_map([
79 ("method", msg_str("auth.logout")),
80 ("params", msg_array(vec![msg_str(&token)])),
81 ]));
82 let _ = send_request(&url, &payload, 5);
83 }
84 self.token = None;
85 }
86
87 pub fn is_authenticated(&self) -> bool {
88 self.token.is_some()
89 }
90
91 pub fn get_workspaces(&self) -> Result<MSFWorkspacesResult> {
92 let token = match self.token.as_ref() {
93 Some(token) => token.clone(),
94 None => {
95 return Ok(MSFWorkspacesResult {
96 workspaces: vec![],
97 error: "Not authenticated".to_string(),
98 });
99 }
100 };
101
102 let url = build_url(&self.host, self.port, self.ssl);
103 let payload = encode(&msg_map([
104 ("method", msg_str("db.workspaces")),
105 ("params", msg_array(vec![msg_str(&token)])),
106 ]));
107
108 match send_request(&url, &payload, 15) {
109 Ok(val) => Ok(MSFWorkspacesResult {
110 workspaces: extract_workspaces(&val),
111 error: String::new(),
112 }),
113 Err(error) => Ok(MSFWorkspacesResult {
114 workspaces: vec![],
115 error,
116 }),
117 }
118 }
119
120 pub fn get_hosts(&self, workspace: Option<&str>) -> Result<MSFHostsResult> {
121 let token = match self.token.as_ref() {
122 Some(token) => token.clone(),
123 None => {
124 return Ok(MSFHostsResult {
125 hosts: vec![],
126 error: "Not authenticated".to_string(),
127 });
128 }
129 };
130
131 let url = build_url(&self.host, self.port, self.ssl);
132 let mut params = vec![msg_str(&token)];
133 if let Some(workspace) = workspace {
134 params.push(MsgValue::Map(BTreeMap::from([(
135 "workspace".to_string(),
136 msg_str(workspace),
137 )])));
138 }
139 let payload = encode(&msg_map([
140 ("method", msg_str("db.hosts")),
141 ("params", msg_array(params)),
142 ]));
143
144 match send_request(&url, &payload, 15) {
145 Ok(val) => Ok(MSFHostsResult {
146 hosts: extract_hosts(&val),
147 error: String::new(),
148 }),
149 Err(error) => Ok(MSFHostsResult {
150 hosts: vec![],
151 error,
152 }),
153 }
154 }
155
156 pub fn check_host_exists(
157 &self,
158 address: &str,
159 workspace: Option<&str>,
160 ) -> Result<MSFHostCheckResult> {
161 let hosts = self.get_hosts(workspace)?;
162 Ok(MSFHostCheckResult {
163 exists: hosts.hosts.iter().any(|host| host.address == address),
164 error: hosts.error,
165 })
166 }
167}
168
169fn build_url(host: &str, port: u16, ssl: bool) -> String {
170 let scheme = if ssl { "https" } else { "http" };
171 format!("{}://{}:{}/api/1.1", scheme, host, port)
172}
173
174fn send_request(
175 url: &str,
176 payload: &[u8],
177 timeout_secs: u64,
178) -> std::result::Result<MsgValue, String> {
179 let mut child = Command::new("curl")
180 .arg("-sS")
181 .arg("--fail")
182 .arg("--max-time")
183 .arg(timeout_secs.to_string())
184 .arg("-H")
185 .arg("Content-Type: binary/message-pack")
186 .arg("--data-binary")
187 .arg("@-")
188 .arg(url)
189 .stdin(Stdio::piped())
190 .stdout(Stdio::piped())
191 .stderr(Stdio::piped())
192 .spawn()
193 .map_err(|e| e.to_string())?;
194
195 if let Some(stdin) = child.stdin.as_mut() {
196 use std::io::Write;
197 stdin.write_all(payload).map_err(|e| e.to_string())?;
198 }
199
200 let output = child.wait_with_output().map_err(|e| e.to_string())?;
201 if !output.status.success() {
202 return Err(String::from_utf8_lossy(&output.stderr).trim().to_string());
203 }
204
205 decode(&output.stdout).map_err(|e| e.to_string())
206}
207
208fn msg_str(value: &str) -> MsgValue {
209 MsgValue::String(value.to_string())
210}
211
212fn msg_array(values: Vec<MsgValue>) -> MsgValue {
213 MsgValue::Array(values)
214}
215
216fn msg_map<const N: usize>(pairs: [(&str, MsgValue); N]) -> MsgValue {
217 let mut map = BTreeMap::new();
218 for (key, value) in pairs {
219 map.insert(key.to_string(), value);
220 }
221 MsgValue::Map(map)
222}
223
224fn extract_workspaces(val: &MsgValue) -> Vec<MSFWorkspaceInfo> {
225 let mut workspaces = Vec::new();
226 if let Some(map) = val.as_map() {
227 if let Some(MsgValue::Array(ws_arr)) = map.get("workspaces") {
228 for ws in ws_arr {
229 if let Some(ws_map) = ws.as_map() {
230 let name = ws_map
231 .get("name")
232 .and_then(MsgValue::as_str)
233 .unwrap_or("")
234 .to_string();
235 let scope = ws_map
236 .get("scope")
237 .and_then(MsgValue::as_str)
238 .unwrap_or("")
239 .to_string();
240 let host_count = ws_map
241 .get("hosts_count")
242 .or_else(|| ws_map.get("host_count"))
243 .and_then(MsgValue::as_u64)
244 .unwrap_or(0);
245
246 workspaces.push(MSFWorkspaceInfo {
247 name,
248 scope,
249 host_count,
250 });
251 }
252 }
253 }
254 }
255 workspaces
256}
257
258fn extract_hosts(val: &MsgValue) -> Vec<MSFHostInfo> {
259 let mut hosts = Vec::new();
260 if let Some(map) = val.as_map() {
261 if let Some(MsgValue::Array(host_arr)) = map.get("hosts") {
262 for host in host_arr {
263 if let Some(host_map) = host.as_map() {
264 let address = host_map
265 .get("address")
266 .and_then(MsgValue::as_str)
267 .unwrap_or("")
268 .to_string();
269 let os_name = host_map
270 .get("os_name")
271 .and_then(MsgValue::as_str)
272 .unwrap_or("")
273 .to_string();
274 let os_flavor = host_map
275 .get("os_flavor")
276 .and_then(MsgValue::as_str)
277 .unwrap_or("")
278 .to_string();
279 let state = host_map
280 .get("state")
281 .and_then(MsgValue::as_str)
282 .unwrap_or("")
283 .to_string();
284 let notes_count = host_map
285 .get("notes_count")
286 .and_then(MsgValue::as_u64)
287 .unwrap_or(0);
288
289 hosts.push(MSFHostInfo {
290 address,
291 os_name,
292 os_flavor,
293 state,
294 notes_count,
295 });
296 }
297 }
298 }
299 }
300 hosts
301}