1use crate::version_text;
2use canic_backup::snapshot::{
3 SnapshotDownloadConfig, SnapshotDownloadError, SnapshotDownloadResult, SnapshotDriver,
4 SnapshotDriverError, SnapshotLifecycleMode,
5};
6use std::{
7 collections::BTreeSet,
8 ffi::OsString,
9 path::{Path, PathBuf},
10 process::Command,
11};
12use thiserror::Error as ThisError;
13
14#[derive(Debug, ThisError)]
19pub enum SnapshotCommandError {
20 #[error("{0}")]
21 Usage(&'static str),
22
23 #[error("missing required option {0}")]
24 MissingOption(&'static str),
25
26 #[error("unknown option {0}")]
27 UnknownOption(String),
28
29 #[error("option {0} requires a value")]
30 MissingValue(&'static str),
31
32 #[error("dfx command failed: {command}\n{stderr}")]
33 DfxFailed { command: String, stderr: String },
34
35 #[error("could not parse snapshot id from dfx output: {0}")]
36 SnapshotIdUnavailable(String),
37
38 #[error(transparent)]
39 Io(#[from] std::io::Error),
40
41 #[error(transparent)]
42 SnapshotDownload(#[from] SnapshotDownloadError),
43}
44
45#[derive(Clone, Debug, Eq, PartialEq)]
50pub struct SnapshotDownloadOptions {
51 pub canister: String,
52 pub out: PathBuf,
53 pub root: Option<String>,
54 pub include_children: bool,
55 pub recursive: bool,
56 pub dry_run: bool,
57 pub lifecycle: SnapshotLifecycleMode,
58 pub network: Option<String>,
59 pub dfx: String,
60}
61
62impl SnapshotDownloadOptions {
63 pub fn parse<I>(args: I) -> Result<Self, SnapshotCommandError>
65 where
66 I: IntoIterator<Item = OsString>,
67 {
68 let mut canister = None;
69 let mut out = None;
70 let mut root = None;
71 let mut include_children = false;
72 let mut recursive = false;
73 let mut dry_run = false;
74 let mut stop_before_snapshot = false;
75 let mut resume_after_snapshot = false;
76 let mut network = None;
77 let mut dfx = "dfx".to_string();
78
79 let mut args = args.into_iter();
80 while let Some(arg) = args.next() {
81 let arg = arg
82 .into_string()
83 .map_err(|_| SnapshotCommandError::Usage(usage()))?;
84 match arg.as_str() {
85 "--canister" => canister = Some(next_value(&mut args, "--canister")?),
86 "--out" => out = Some(PathBuf::from(next_value(&mut args, "--out")?)),
87 "--root" => root = Some(next_value(&mut args, "--root")?),
88 "--include-children" => include_children = true,
89 "--recursive" => {
90 recursive = true;
91 include_children = true;
92 }
93 "--dry-run" => dry_run = true,
94 "--stop-before-snapshot" => stop_before_snapshot = true,
95 "--resume-after-snapshot" => resume_after_snapshot = true,
96 "--network" => network = Some(next_value(&mut args, "--network")?),
97 "--dfx" => dfx = next_value(&mut args, "--dfx")?,
98 "--help" | "-h" => return Err(SnapshotCommandError::Usage(usage())),
99 _ => return Err(SnapshotCommandError::UnknownOption(arg)),
100 }
101 }
102
103 Ok(Self {
104 canister: canister.ok_or(SnapshotCommandError::MissingOption("--canister"))?,
105 out: out.ok_or(SnapshotCommandError::MissingOption("--out"))?,
106 root,
107 include_children,
108 recursive,
109 dry_run,
110 lifecycle: SnapshotLifecycleMode::from_flags(
111 stop_before_snapshot,
112 resume_after_snapshot,
113 ),
114 network,
115 dfx,
116 })
117 }
118}
119
120pub fn run<I>(args: I) -> Result<(), SnapshotCommandError>
122where
123 I: IntoIterator<Item = OsString>,
124{
125 let mut args = args.into_iter();
126 let Some(command) = args.next().and_then(|arg| arg.into_string().ok()) else {
127 return Err(SnapshotCommandError::Usage(usage()));
128 };
129
130 match command.as_str() {
131 "download" => {
132 let options = SnapshotDownloadOptions::parse(args)?;
133 let result = download_snapshots(&options)?;
134 for command in result.planned_commands {
135 println!("{command}");
136 }
137 for artifact in result.artifacts {
138 println!(
139 "{} {} {}",
140 artifact.canister_id,
141 artifact.snapshot_id,
142 artifact.path.display()
143 );
144 }
145 Ok(())
146 }
147 "help" | "--help" | "-h" => {
148 println!("{}", usage());
149 Ok(())
150 }
151 "version" | "--version" | "-V" => {
152 println!("{}", version_text());
153 Ok(())
154 }
155 _ => Err(SnapshotCommandError::UnknownOption(command)),
156 }
157}
158
159pub fn download_snapshots(
161 options: &SnapshotDownloadOptions,
162) -> Result<SnapshotDownloadResult, SnapshotCommandError> {
163 let config = SnapshotDownloadConfig {
164 canister: options.canister.clone(),
165 out: options.out.clone(),
166 root: options.root.clone(),
167 include_children: options.include_children,
168 recursive: options.recursive,
169 dry_run: options.dry_run,
170 lifecycle: options.lifecycle,
171 backup_id: backup_id(options),
172 created_at: timestamp_placeholder(),
173 tool_name: "canic-cli".to_string(),
174 tool_version: env!("CARGO_PKG_VERSION").to_string(),
175 environment: options
176 .network
177 .clone()
178 .unwrap_or_else(|| "local".to_string()),
179 };
180 let mut driver = DfxSnapshotDriver { options };
181 canic_backup::snapshot::download_snapshots(&config, &mut driver)
182 .map_err(SnapshotCommandError::from)
183}
184
185struct DfxSnapshotDriver<'a> {
190 options: &'a SnapshotDownloadOptions,
191}
192
193impl SnapshotDriver for DfxSnapshotDriver<'_> {
194 fn registry_json(&mut self, root: &str) -> Result<String, SnapshotDriverError> {
196 call_subnet_registry(self.options, root).map_err(driver_error)
197 }
198
199 fn create_snapshot(&mut self, canister_id: &str) -> Result<String, SnapshotDriverError> {
201 create_snapshot(self.options, canister_id).map_err(driver_error)
202 }
203
204 fn stop_canister(&mut self, canister_id: &str) -> Result<(), SnapshotDriverError> {
206 stop_canister(self.options, canister_id).map_err(driver_error)
207 }
208
209 fn start_canister(&mut self, canister_id: &str) -> Result<(), SnapshotDriverError> {
211 start_canister(self.options, canister_id).map_err(driver_error)
212 }
213
214 fn download_snapshot(
216 &mut self,
217 canister_id: &str,
218 snapshot_id: &str,
219 artifact_path: &Path,
220 ) -> Result<(), SnapshotDriverError> {
221 download_snapshot(self.options, canister_id, snapshot_id, artifact_path)
222 .map_err(driver_error)
223 }
224
225 fn create_snapshot_command(&self, canister_id: &str) -> String {
227 create_snapshot_command_display(self.options, canister_id)
228 }
229
230 fn stop_canister_command(&self, canister_id: &str) -> String {
232 stop_canister_command_display(self.options, canister_id)
233 }
234
235 fn start_canister_command(&self, canister_id: &str) -> String {
237 start_canister_command_display(self.options, canister_id)
238 }
239
240 fn download_snapshot_command(
242 &self,
243 canister_id: &str,
244 snapshot_id: &str,
245 artifact_path: &Path,
246 ) -> String {
247 download_snapshot_command_display(self.options, canister_id, snapshot_id, artifact_path)
248 }
249}
250
251fn driver_error(error: SnapshotCommandError) -> SnapshotDriverError {
253 Box::new(error)
254}
255
256fn call_subnet_registry(
258 options: &SnapshotDownloadOptions,
259 root: &str,
260) -> Result<String, SnapshotCommandError> {
261 let mut command = Command::new(&options.dfx);
262 command.arg("canister");
263 add_canister_network_args(&mut command, options);
264 command.args(["call", root, "canic_subnet_registry", "--output", "json"]);
265 run_output(&mut command)
266}
267
268fn create_snapshot(
270 options: &SnapshotDownloadOptions,
271 canister_id: &str,
272) -> Result<String, SnapshotCommandError> {
273 let before = list_snapshot_ids(options, canister_id)?;
274 let mut command = Command::new(&options.dfx);
275 command.arg("canister");
276 add_canister_network_args(&mut command, options);
277 command.args(["snapshot", "create", canister_id]);
278 let output = run_output_with_stderr(&mut command)?;
279 if let Some(snapshot_id) = parse_snapshot_id(&output) {
280 return Ok(snapshot_id);
281 }
282
283 let before = before.into_iter().collect::<BTreeSet<_>>();
284 let mut new_ids = list_snapshot_ids(options, canister_id)?
285 .into_iter()
286 .filter(|snapshot_id| !before.contains(snapshot_id))
287 .collect::<Vec<_>>();
288 if new_ids.len() == 1 {
289 Ok(new_ids.remove(0))
290 } else {
291 Err(SnapshotCommandError::SnapshotIdUnavailable(output))
292 }
293}
294
295fn list_snapshot_ids(
297 options: &SnapshotDownloadOptions,
298 canister_id: &str,
299) -> Result<Vec<String>, SnapshotCommandError> {
300 let mut command = Command::new(&options.dfx);
301 command.arg("canister");
302 add_canister_network_args(&mut command, options);
303 command.args(["snapshot", "list", canister_id]);
304 let output = run_output(&mut command)?;
305 Ok(parse_snapshot_list_ids(&output))
306}
307
308fn stop_canister(
310 options: &SnapshotDownloadOptions,
311 canister_id: &str,
312) -> Result<(), SnapshotCommandError> {
313 let mut command = Command::new(&options.dfx);
314 command.arg("canister");
315 add_canister_network_args(&mut command, options);
316 command.args(["stop", canister_id]);
317 run_status(&mut command)
318}
319
320fn start_canister(
322 options: &SnapshotDownloadOptions,
323 canister_id: &str,
324) -> Result<(), SnapshotCommandError> {
325 let mut command = Command::new(&options.dfx);
326 command.arg("canister");
327 add_canister_network_args(&mut command, options);
328 command.args(["start", canister_id]);
329 run_status(&mut command)
330}
331
332fn download_snapshot(
334 options: &SnapshotDownloadOptions,
335 canister_id: &str,
336 snapshot_id: &str,
337 artifact_path: &Path,
338) -> Result<(), SnapshotCommandError> {
339 let mut command = Command::new(&options.dfx);
340 command.arg("canister");
341 add_canister_network_args(&mut command, options);
342 command.args(["snapshot", "download", canister_id, snapshot_id, "--dir"]);
343 command.arg(artifact_path);
344 run_status(&mut command)
345}
346
347fn add_canister_network_args(command: &mut Command, options: &SnapshotDownloadOptions) {
349 if let Some(network) = &options.network {
350 command.args(["--network", network]);
351 }
352}
353
354fn run_output(command: &mut Command) -> Result<String, SnapshotCommandError> {
356 let display = command_display(command);
357 let output = command.output()?;
358 if output.status.success() {
359 Ok(String::from_utf8_lossy(&output.stdout).trim().to_string())
360 } else {
361 Err(SnapshotCommandError::DfxFailed {
362 command: display,
363 stderr: String::from_utf8_lossy(&output.stderr).to_string(),
364 })
365 }
366}
367
368fn run_output_with_stderr(command: &mut Command) -> Result<String, SnapshotCommandError> {
370 let display = command_display(command);
371 let output = command.output()?;
372 if output.status.success() {
373 let mut text = String::from_utf8_lossy(&output.stdout).to_string();
374 text.push_str(&String::from_utf8_lossy(&output.stderr));
375 Ok(text.trim().to_string())
376 } else {
377 Err(SnapshotCommandError::DfxFailed {
378 command: display,
379 stderr: String::from_utf8_lossy(&output.stderr).to_string(),
380 })
381 }
382}
383
384fn run_status(command: &mut Command) -> Result<(), SnapshotCommandError> {
386 let display = command_display(command);
387 let output = command.output()?;
388 if output.status.success() {
389 Ok(())
390 } else {
391 Err(SnapshotCommandError::DfxFailed {
392 command: display,
393 stderr: String::from_utf8_lossy(&output.stderr).to_string(),
394 })
395 }
396}
397
398fn command_display(command: &Command) -> String {
400 let mut parts = vec![command.get_program().to_string_lossy().to_string()];
401 parts.extend(
402 command
403 .get_args()
404 .map(|arg| arg.to_string_lossy().to_string()),
405 );
406 parts.join(" ")
407}
408
409fn create_snapshot_command_display(options: &SnapshotDownloadOptions, canister_id: &str) -> String {
411 let mut command = Command::new(&options.dfx);
412 command.arg("canister");
413 add_canister_network_args(&mut command, options);
414 command.args(["snapshot", "create", canister_id]);
415 command_display(&command)
416}
417
418fn download_snapshot_command_display(
420 options: &SnapshotDownloadOptions,
421 canister_id: &str,
422 snapshot_id: &str,
423 artifact_path: &Path,
424) -> String {
425 let mut command = Command::new(&options.dfx);
426 command.arg("canister");
427 add_canister_network_args(&mut command, options);
428 command.args(["snapshot", "download", canister_id, snapshot_id, "--dir"]);
429 command.arg(artifact_path);
430 command_display(&command)
431}
432
433fn stop_canister_command_display(options: &SnapshotDownloadOptions, canister_id: &str) -> String {
435 let mut command = Command::new(&options.dfx);
436 command.arg("canister");
437 add_canister_network_args(&mut command, options);
438 command.args(["stop", canister_id]);
439 command_display(&command)
440}
441
442fn start_canister_command_display(options: &SnapshotDownloadOptions, canister_id: &str) -> String {
444 let mut command = Command::new(&options.dfx);
445 command.arg("canister");
446 add_canister_network_args(&mut command, options);
447 command.args(["start", canister_id]);
448 command_display(&command)
449}
450
451fn parse_snapshot_id(output: &str) -> Option<String> {
453 output
454 .split(|c: char| c.is_whitespace() || matches!(c, '"' | '\'' | ':' | ','))
455 .filter(|part| !part.is_empty())
456 .rev()
457 .find(|part| {
458 part.chars()
459 .all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.'))
460 })
461 .map(str::to_string)
462}
463
464fn parse_snapshot_list_ids(output: &str) -> Vec<String> {
466 output
467 .lines()
468 .filter_map(|line| {
469 line.split_once(':')
470 .map(|(snapshot_id, _)| snapshot_id.trim())
471 })
472 .filter(|snapshot_id| !snapshot_id.is_empty())
473 .map(str::to_string)
474 .collect()
475}
476
477fn backup_id(options: &SnapshotDownloadOptions) -> String {
479 options
480 .out
481 .file_name()
482 .and_then(|name| name.to_str())
483 .map_or_else(|| "snapshot-download".to_string(), str::to_string)
484}
485
486fn timestamp_placeholder() -> String {
488 "unknown".to_string()
489}
490
491fn next_value<I>(args: &mut I, option: &'static str) -> Result<String, SnapshotCommandError>
493where
494 I: Iterator<Item = OsString>,
495{
496 args.next()
497 .and_then(|value| value.into_string().ok())
498 .ok_or(SnapshotCommandError::MissingValue(option))
499}
500
501const fn usage() -> &'static str {
503 "usage: canic snapshot download --canister <id> --out <dir> [--root <id>] [--include-children] [--recursive] [--dry-run] [--stop-before-snapshot] [--resume-after-snapshot] [--network <name>]"
504}
505
506#[cfg(test)]
507mod tests {
508 use super::*;
509
510 const ROOT: &str = "aaaaa-aa";
511
512 #[test]
514 fn parses_snapshot_id_from_output() {
515 let snapshot_id = parse_snapshot_id("Created snapshot: snap_abc-123\n");
516
517 assert_eq!(snapshot_id.as_deref(), Some("snap_abc-123"));
518 }
519
520 #[test]
522 fn parses_snapshot_ids_from_list_output() {
523 let snapshot_ids = parse_snapshot_list_ids(
524 "0000000000000000ffffffffff9000050101: 213.76 MiB, taken at 2026-05-03 12:20:53 UTC\n",
525 );
526
527 assert_eq!(snapshot_ids, vec!["0000000000000000ffffffffff9000050101"]);
528 }
529
530 #[test]
532 fn parses_download_options() {
533 let options = SnapshotDownloadOptions::parse([
534 OsString::from("--canister"),
535 OsString::from(ROOT),
536 OsString::from("--out"),
537 OsString::from("backups/test"),
538 OsString::from("--root"),
539 OsString::from(ROOT),
540 OsString::from("--recursive"),
541 OsString::from("--dry-run"),
542 OsString::from("--stop-before-snapshot"),
543 OsString::from("--resume-after-snapshot"),
544 ])
545 .expect("parse options");
546
547 assert_eq!(options.canister, ROOT);
548 assert!(options.include_children);
549 assert!(options.recursive);
550 assert!(options.dry_run);
551 assert_eq!(options.root.as_deref(), Some(ROOT));
552 assert_eq!(options.lifecycle, SnapshotLifecycleMode::StopAndResume);
553 }
554}