atlas_cli/cli/
handlers.rs

1use crate::error::{Error, Result};
2
3use super::commands::{
4    CCAttestationCommands, DatasetCommands, EvaluationCommands, ManifestCommands, ModelCommands,
5    PipelineCommands, SoftwareCommands,
6};
7use crate::cc_attestation;
8use crate::manifest;
9use crate::manifest::config::ManifestCreationConfig;
10use crate::manifest::dataset::list_dataset_manifests;
11use crate::slsa;
12use crate::storage::database::DatabaseStorage;
13use crate::storage::filesystem::FilesystemStorage;
14use crate::storage::rekor::RekorStorage;
15
16use crate::StorageBackend;
17
18pub fn handle_dataset_command(cmd: DatasetCommands) -> Result<()> {
19    let _storage = RekorStorage::new()?;
20    match cmd {
21        DatasetCommands::Create {
22            paths,
23            ingredient_names,
24            name,
25            author_org,
26            author_name,
27            description,
28            linked_manifests,
29            storage_type,
30            storage_url,
31            print,
32            encoding,
33            key,
34            hash_alg,
35            with_tdx,
36        } => {
37            let storage: Option<&'static dyn StorageBackend> = match storage_type.as_str() {
38                "database" => {
39                    let db_storage = Box::new(DatabaseStorage::new(*storage_url.clone())?);
40                    Some(Box::leak(db_storage))
41                }
42                "rekor" => {
43                    let rekor_storage = Box::new(RekorStorage::new_with_url(*storage_url.clone())?);
44                    Some(Box::leak(rekor_storage))
45                }
46                "local-fs" => {
47                    let fs_storage = Box::new(FilesystemStorage::new(storage_url.as_str())?);
48                    Some(Box::leak(fs_storage))
49                }
50                _ => None,
51            };
52
53            let config = ManifestCreationConfig {
54                paths,
55                ingredient_names,
56                name,
57                author_org,
58                author_name,
59                description,
60                linked_manifests,
61                storage,
62                print,
63                output_encoding: encoding,
64                key_path: key,
65                hash_alg: hash_alg.to_cose_algorithm(),
66                with_cc: with_tdx,
67                software_type: None,
68                version: None,
69                custom_fields: None,
70            };
71
72            manifest::create_dataset_manifest(config)
73        }
74        DatasetCommands::List {
75            storage_type,
76            storage_url,
77        } => {
78            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
79                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
80                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
81                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
82                _ => return Err(Error::Validation("Invalid storage type".to_string())),
83            };
84
85            list_dataset_manifests(storage.as_ref())
86        }
87        DatasetCommands::Verify {
88            id,
89            storage_type,
90            storage_url,
91        } => {
92            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
93                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
94                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
95                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
96                _ => return Err(Error::Validation("Invalid storage type".to_string())),
97            };
98
99            manifest::verify_dataset_manifest(&id, storage.as_ref())
100        }
101    }
102}
103
104pub fn handle_model_command(cmd: ModelCommands) -> Result<()> {
105    let _storage = RekorStorage::new()?;
106    match cmd {
107        ModelCommands::Create {
108            paths,
109            ingredient_names,
110            name,
111            author_org,
112            author_name,
113            description,
114            linked_manifests,
115            storage_type,
116            storage_url,
117            print,
118            encoding,
119            format,
120            key,
121            hash_alg,
122            with_tdx,
123        } => {
124            let storage: Option<&'static dyn StorageBackend> = match storage_type.as_str() {
125                "database" => {
126                    let db_storage = Box::new(DatabaseStorage::new(*storage_url.clone())?);
127                    Some(Box::leak(db_storage))
128                }
129                "rekor" => {
130                    let rekor_storage = Box::new(RekorStorage::new_with_url(*storage_url.clone())?);
131                    Some(Box::leak(rekor_storage))
132                }
133                "local-fs" => {
134                    let fs_storage = Box::new(FilesystemStorage::new(storage_url.as_str())?);
135                    Some(Box::leak(fs_storage))
136                }
137                _ => None,
138            };
139
140            let config = ManifestCreationConfig {
141                paths,
142                ingredient_names,
143                name,
144                author_org,
145                author_name,
146                description,
147                linked_manifests,
148                storage,
149                print,
150                output_encoding: encoding,
151                key_path: key,
152                hash_alg: hash_alg.to_cose_algorithm(),
153                with_cc: with_tdx,
154                software_type: None,
155                version: None,
156                custom_fields: None,
157            };
158
159            match format.as_str() {
160                "standalone" => manifest::create_model_manifest(config),
161                "oms" => manifest::common::create_oms_manifest(config),
162                _ => {
163                    return Err(Error::InitializationError(
164                        "Unsupported output format".to_string(),
165                    ));
166                }
167            }
168        }
169        ModelCommands::List {
170            storage_type,
171            storage_url,
172        } => {
173            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
174                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
175                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
176                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
177                _ => return Err(Error::Validation("Invalid storage type".to_string())),
178            };
179
180            manifest::list_model_manifest(storage.as_ref())
181        }
182        ModelCommands::Verify {
183            id,
184            storage_type,
185            storage_url,
186        } => {
187            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
188                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
189                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
190                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
191                _ => return Err(Error::Validation("Invalid storage type".to_string())),
192            };
193
194            manifest::verify_model_manifest(&id, storage.as_ref())
195        }
196        ModelCommands::LinkDataset {
197            model_id,
198            dataset_id,
199            storage_type,
200            storage_url,
201        } => {
202            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
203                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
204                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
205                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
206                _ => return Err(Error::Validation("Invalid storage type".to_string())),
207            };
208
209            let updated_manifest =
210                manifest::linking::link_dataset_to_model(&model_id, &dataset_id, storage.as_ref())?;
211
212            println!("Successfully linked dataset {dataset_id} to model {model_id}");
213            println!("Updated manifest ID: {}", updated_manifest.instance_id);
214
215            Ok(())
216        }
217    }
218}
219
220pub fn handle_manifest_command(cmd: ManifestCommands) -> Result<()> {
221    match cmd {
222        ManifestCommands::Link {
223            source,
224            target,
225            storage_type,
226            storage_url,
227        } => {
228            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
229                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
230                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
231                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
232                _ => return Err(Error::Validation("Invalid storage type".to_string())),
233            };
234
235            manifest::link_manifests(&source, &target, &*storage)
236        }
237        ManifestCommands::Show {
238            id,
239            storage_type,
240            storage_url,
241        } => {
242            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
243                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
244                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
245                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
246                _ => return Err(Error::Validation("Invalid storage type".to_string())),
247            };
248
249            manifest::show_manifest(&id, &*storage)
250        }
251        ManifestCommands::Validate {
252            id,
253            storage_type,
254            storage_url,
255        } => {
256            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
257                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
258                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
259                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
260                _ => return Err(Error::Validation("Invalid storage type".to_string())),
261            };
262
263            manifest::validate_linked_manifests(&id, &*storage)
264        }
265        ManifestCommands::VerifyLink {
266            source,
267            target,
268            storage_type,
269            storage_url,
270        } => {
271            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
272                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
273                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
274                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
275                _ => return Err(Error::Validation("Invalid storage type".to_string())),
276            };
277
278            let result = manifest::verify_manifest_link(&source, &target, &*storage)?;
279            if result {
280                println!("Link verification successful");
281                Ok(())
282            } else {
283                Err(Error::Validation("Link verification failed".to_string()))
284            }
285        }
286        ManifestCommands::Export {
287            id,
288            storage_type,
289            storage_url,
290            encoding,
291            output,
292            max_depth,
293        } => {
294            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
295                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
296                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
297                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
298                _ => return Err(Error::Validation("Invalid storage type".to_string())),
299            };
300
301            manifest::export_provenance(
302                &id,
303                &*storage,
304                encoding.as_str(),
305                output.as_deref(),
306                max_depth,
307            )
308        }
309    }
310}
311
312pub fn handle_evaluation_command(cmd: EvaluationCommands) -> Result<()> {
313    match cmd {
314        EvaluationCommands::Create {
315            path,
316            name,
317            model_id,
318            dataset_id,
319            metrics,
320            author_org,
321            author_name,
322            description,
323            storage_type,
324            storage_url,
325            print,
326            encoding,
327            key,
328            hash_alg,
329        } => {
330            let storage: Option<&'static dyn StorageBackend> = match storage_type.as_str() {
331                "database" => {
332                    let db_storage = Box::new(DatabaseStorage::new(*storage_url.clone())?);
333                    Some(Box::leak(db_storage))
334                }
335                "rekor" => {
336                    let rekor_storage = Box::new(RekorStorage::new_with_url(*storage_url.clone())?);
337                    Some(Box::leak(rekor_storage))
338                }
339                "local-fs" => {
340                    let fs_storage = Box::new(FilesystemStorage::new(storage_url.as_str())?);
341                    Some(Box::leak(fs_storage))
342                }
343                _ => None,
344            };
345
346            let config = ManifestCreationConfig {
347                paths: vec![path],
348                ingredient_names: vec!["Evaluation Results".to_string()],
349                name,
350                author_org,
351                author_name,
352                description,
353                linked_manifests: None, // Will be populated by create_manifest
354                storage,
355                print,
356                output_encoding: encoding,
357                key_path: key,
358                hash_alg: hash_alg.to_cose_algorithm(),
359                with_cc: false,
360                software_type: None,
361                version: None,
362                custom_fields: None, // Will be populated by create_manifest
363            };
364
365            manifest::evaluation::create_manifest(config, model_id, dataset_id, metrics)
366        }
367        EvaluationCommands::List {
368            storage_type,
369            storage_url,
370        } => {
371            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
372                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
373                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
374                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
375                _ => return Err(Error::Validation("Invalid storage type".to_string())),
376            };
377
378            manifest::evaluation::list_evaluation_manifests(storage.as_ref())
379        }
380        EvaluationCommands::Verify {
381            id,
382            storage_type,
383            storage_url,
384        } => {
385            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
386                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
387                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
388                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
389                _ => return Err(Error::Validation("Invalid storage type".to_string())),
390            };
391
392            manifest::evaluation::verify_evaluation_manifest(&id, storage.as_ref())
393        }
394    }
395}
396
397pub fn handle_cc_attestation_command(cmd: CCAttestationCommands) -> Result<()> {
398    match cmd {
399        CCAttestationCommands::Show => {
400            let _r = cc_attestation::get_report(true).unwrap();
401            Ok(())
402        }
403
404        CCAttestationCommands::GetLaunchMeasurement => {
405            let m = cc_attestation::get_launch_measurement().unwrap();
406            println!("Launch measurement raw bytes: 0x{}", hex::encode(m));
407            Ok(())
408        }
409
410        CCAttestationCommands::VerifyLaunch { host_platform } => {
411            let result = cc_attestation::verify_launch_endorsement(&host_platform).unwrap();
412            if result {
413                println!(
414                    "Passed: launch endorsement verification for {host_platform} host platform"
415                );
416            } else {
417                println!(
418                    "Failed: launch endorsement verification for {host_platform} host platform"
419                );
420            }
421            Ok(())
422        }
423    }
424}
425
426pub fn handle_software_command(cmd: SoftwareCommands) -> Result<()> {
427    match cmd {
428        SoftwareCommands::Create {
429            paths,
430            ingredient_names,
431            name,
432            software_type,
433            version,
434            author_org,
435            author_name,
436            description,
437            linked_manifests,
438            storage_type,
439            storage_url,
440            print,
441            encoding,
442            key,
443            hash_alg,
444            with_tdx,
445        } => {
446            let storage: Option<&'static dyn StorageBackend> = match storage_type.as_str() {
447                "database" => {
448                    let db_storage = Box::new(DatabaseStorage::new(*storage_url.clone())?);
449                    Some(Box::leak(db_storage))
450                }
451                "rekor" => {
452                    let rekor_storage = Box::new(RekorStorage::new_with_url(*storage_url.clone())?);
453                    Some(Box::leak(rekor_storage))
454                }
455                "local-fs" => {
456                    let fs_storage = Box::new(FilesystemStorage::new(storage_url.as_str())?);
457                    Some(Box::leak(fs_storage))
458                }
459                _ => None,
460            };
461
462            let config = ManifestCreationConfig {
463                paths,
464                ingredient_names,
465                name,
466                author_org,
467                author_name,
468                description,
469                linked_manifests,
470                storage,
471                print,
472                output_encoding: encoding,
473                key_path: key,
474                hash_alg: hash_alg.to_cose_algorithm(),
475                with_cc: with_tdx,
476                software_type: Some(software_type.clone()),
477                version: version.clone(),
478                custom_fields: None,
479            };
480
481            manifest::software::create_manifest(config, software_type, version)
482        }
483        SoftwareCommands::List {
484            storage_type,
485            storage_url,
486        } => {
487            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
488                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
489                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
490                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
491                _ => return Err(Error::Validation("Invalid storage type".to_string())),
492            };
493
494            manifest::software::list_software_manifests(storage.as_ref())
495        }
496        SoftwareCommands::Verify {
497            id,
498            storage_type,
499            storage_url,
500        } => {
501            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
502                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
503                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
504                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
505                _ => return Err(Error::Validation("Invalid storage type".to_string())),
506            };
507
508            manifest::software::verify_software_manifest(&id, storage.as_ref())
509        }
510        SoftwareCommands::LinkModel {
511            software_id,
512            model_id,
513            storage_type,
514            storage_url,
515        } => {
516            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
517                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
518                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
519                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
520                _ => return Err(Error::Validation("Invalid storage type".to_string())),
521            };
522
523            // Link software to model
524            manifest::link_manifests(&model_id, &software_id, storage.as_ref())
525        }
526        SoftwareCommands::LinkDataset {
527            software_id,
528            dataset_id,
529            storage_type,
530            storage_url,
531        } => {
532            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
533                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
534                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
535                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
536                _ => return Err(Error::Validation("Invalid storage type".to_string())),
537            };
538
539            // Link software to dataset
540            manifest::link_manifests(&dataset_id, &software_id, storage.as_ref())
541        }
542    }
543}
544
545pub fn handle_pipeline_command(cmd: PipelineCommands) -> Result<()> {
546    match cmd {
547        PipelineCommands::GenerateProvenance {
548            inputs,
549            pipeline,
550            products,
551            key,
552            hash_alg,
553            encoding,
554            print,
555            storage_type,
556            storage_url,
557            with_tdx,
558        } => {
559            let storage: Option<&'static dyn StorageBackend> = match storage_type.as_str() {
560                "local-fs" => {
561                    let fs_storage = Box::new(FilesystemStorage::new(storage_url.as_str())?);
562                    Some(Box::leak(fs_storage))
563                }
564                _ => None,
565            };
566
567            slsa::cli::generate_build_provenance(
568                inputs,
569                pipeline,
570                products,
571                key,
572                hash_alg.to_cose_algorithm(),
573                encoding,
574                print,
575                storage,
576                with_tdx,
577            )
578        }
579    }
580}
atlas_cli/cli/handlers.rs

atlas_cli/cli/
handlers.rs
