Skip to main content

atlas_cli/cli/
handlers.rs

1use crate::error::{Error, Result};
2
3use super::commands::{
4    CCAttestationCommands, DatasetCommands, EvaluationCommands, ManifestCommands, ModelCommands,
5    PipelineCommands, RekorCommands, 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;
17use std::path::PathBuf;
18
19fn build_rekor_storage(
20    url: String,
21    key: &Option<PathBuf>,
22    cert: &Option<PathBuf>,
23    fulcio: bool,
24    oidc_token: &Option<String>,
25) -> Result<RekorStorage> {
26    let mut storage = RekorStorage::new_with_url(url)?.with_key(key.clone());
27    if fulcio {
28        let token = oidc_token.as_ref().ok_or_else(|| {
29            Error::Validation("--oidc-token is required when --fulcio is set".to_string())
30        })?;
31        storage = storage.with_fulcio(token.clone());
32    } else {
33        storage = storage.with_cert(cert.clone());
34    }
35    Ok(storage)
36}
37
38pub fn handle_dataset_command(cmd: DatasetCommands) -> Result<()> {
39    let _storage = RekorStorage::new()?;
40    match cmd {
41        DatasetCommands::Create {
42            paths,
43            ingredient_names,
44            name,
45            author_org,
46            author_name,
47            description,
48            linked_manifests,
49            storage_type,
50            storage_url,
51            print,
52            encoding,
53            key,
54            cert,
55            fulcio,
56            oidc_token,
57            hash_alg,
58            with_tdx,
59        } => {
60            let storage: Option<&'static dyn StorageBackend> = match storage_type.as_str() {
61                "database" => {
62                    let db_storage = Box::new(DatabaseStorage::new(*storage_url.clone())?);
63                    Some(Box::leak(db_storage))
64                }
65                "rekor" => {
66                    let rekor_storage = Box::new(build_rekor_storage(
67                        *storage_url.clone(),
68                        &key,
69                        &cert,
70                        fulcio,
71                        &oidc_token,
72                    )?);
73                    Some(Box::leak(rekor_storage))
74                }
75                "local-fs" => {
76                    let fs_storage = Box::new(FilesystemStorage::new(storage_url.as_str())?);
77                    Some(Box::leak(fs_storage))
78                }
79                _ => None,
80            };
81
82            let config = ManifestCreationConfig {
83                paths,
84                ingredient_names,
85                name,
86                author_org,
87                author_name,
88                description,
89                linked_manifests,
90                storage,
91                print,
92                output_encoding: encoding,
93                key_path: key,
94                hash_alg: hash_alg.to_cose_algorithm(),
95                with_cc: with_tdx,
96                software_type: None,
97                version: None,
98                custom_fields: None,
99            };
100
101            manifest::create_dataset_manifest(config)
102        }
103        DatasetCommands::List {
104            storage_type,
105            storage_url,
106        } => {
107            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
108                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
109                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
110                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
111                _ => return Err(Error::Validation("Invalid storage type".to_string())),
112            };
113
114            list_dataset_manifests(storage.as_ref())
115        }
116        DatasetCommands::Verify {
117            id,
118            storage_type,
119            storage_url,
120        } => {
121            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
122                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
123                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
124                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
125                _ => return Err(Error::Validation("Invalid storage type".to_string())),
126            };
127
128            manifest::verify_dataset_manifest(&id, storage.as_ref())
129        }
130    }
131}
132
133pub fn handle_model_command(cmd: ModelCommands) -> Result<()> {
134    let _storage = RekorStorage::new()?;
135    match cmd {
136        ModelCommands::Create {
137            paths,
138            ingredient_names,
139            name,
140            author_org,
141            author_name,
142            description,
143            linked_manifests,
144            storage_type,
145            storage_url,
146            print,
147            encoding,
148            format,
149            key,
150            cert,
151            fulcio,
152            oidc_token,
153            hash_alg,
154            with_tdx,
155        } => {
156            let storage: Option<&'static dyn StorageBackend> = match storage_type.as_str() {
157                "database" => {
158                    let db_storage = Box::new(DatabaseStorage::new(*storage_url.clone())?);
159                    Some(Box::leak(db_storage))
160                }
161                "rekor" => {
162                    let rekor_storage = Box::new(build_rekor_storage(
163                        *storage_url.clone(),
164                        &key,
165                        &cert,
166                        fulcio,
167                        &oidc_token,
168                    )?);
169                    Some(Box::leak(rekor_storage))
170                }
171                "local-fs" => {
172                    let fs_storage = Box::new(FilesystemStorage::new(storage_url.as_str())?);
173                    Some(Box::leak(fs_storage))
174                }
175                _ => None,
176            };
177
178            let config = ManifestCreationConfig {
179                paths,
180                ingredient_names,
181                name,
182                author_org,
183                author_name,
184                description,
185                linked_manifests,
186                storage,
187                print,
188                output_encoding: encoding,
189                key_path: key,
190                hash_alg: hash_alg.to_cose_algorithm(),
191                with_cc: with_tdx,
192                software_type: None,
193                version: None,
194                custom_fields: None,
195            };
196
197            match format.as_str() {
198                "standalone" => manifest::create_model_manifest(config),
199                "oms" => manifest::common::create_oms_manifest(config),
200                _ => {
201                    return Err(Error::InitializationError(
202                        "Unsupported output format".to_string(),
203                    ));
204                }
205            }
206        }
207        ModelCommands::List {
208            storage_type,
209            storage_url,
210        } => {
211            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
212                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
213                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
214                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
215                _ => return Err(Error::Validation("Invalid storage type".to_string())),
216            };
217
218            manifest::list_model_manifest(storage.as_ref())
219        }
220        ModelCommands::Verify {
221            id,
222            storage_type,
223            storage_url,
224        } => {
225            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
226                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
227                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
228                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
229                _ => return Err(Error::Validation("Invalid storage type".to_string())),
230            };
231
232            manifest::verify_model_manifest(&id, storage.as_ref())
233        }
234        ModelCommands::LinkDataset {
235            model_id,
236            dataset_id,
237            storage_type,
238            storage_url,
239        } => {
240            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
241                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
242                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
243                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
244                _ => return Err(Error::Validation("Invalid storage type".to_string())),
245            };
246
247            let updated_manifest =
248                manifest::linking::link_dataset_to_model(&model_id, &dataset_id, storage.as_ref())?;
249
250            println!("Successfully linked dataset {dataset_id} to model {model_id}");
251            println!("Updated manifest ID: {}", updated_manifest.instance_id);
252
253            Ok(())
254        }
255    }
256}
257
258pub fn handle_manifest_command(cmd: ManifestCommands) -> Result<()> {
259    match cmd {
260        ManifestCommands::Link {
261            source,
262            target,
263            storage_type,
264            storage_url,
265        } => {
266            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
267                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
268                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
269                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
270                _ => return Err(Error::Validation("Invalid storage type".to_string())),
271            };
272
273            manifest::link_manifests(&source, &target, &*storage)
274        }
275        ManifestCommands::Show {
276            id,
277            storage_type,
278            storage_url,
279        } => {
280            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
281                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
282                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
283                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
284                _ => return Err(Error::Validation("Invalid storage type".to_string())),
285            };
286
287            manifest::show_manifest(&id, &*storage)
288        }
289        ManifestCommands::Validate {
290            id,
291            storage_type,
292            storage_url,
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::validate_linked_manifests(&id, &*storage)
302        }
303        ManifestCommands::VerifyLink {
304            source,
305            target,
306            storage_type,
307            storage_url,
308        } => {
309            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
310                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
311                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
312                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
313                _ => return Err(Error::Validation("Invalid storage type".to_string())),
314            };
315
316            let result = manifest::verify_manifest_link(&source, &target, &*storage)?;
317            if result {
318                println!("Link verification successful");
319                Ok(())
320            } else {
321                Err(Error::Validation("Link verification failed".to_string()))
322            }
323        }
324        ManifestCommands::Export {
325            id,
326            storage_type,
327            storage_url,
328            encoding,
329            output,
330            max_depth,
331        } => {
332            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
333                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
334                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
335                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
336                _ => return Err(Error::Validation("Invalid storage type".to_string())),
337            };
338
339            manifest::export_provenance(
340                &id,
341                &*storage,
342                encoding.as_str(),
343                output.as_deref(),
344                max_depth,
345            )
346        }
347    }
348}
349
350pub fn handle_evaluation_command(cmd: EvaluationCommands) -> Result<()> {
351    match cmd {
352        EvaluationCommands::Create {
353            path,
354            name,
355            model_id,
356            dataset_id,
357            metrics,
358            author_org,
359            author_name,
360            description,
361            storage_type,
362            storage_url,
363            print,
364            encoding,
365            key,
366            cert,
367            fulcio,
368            oidc_token,
369            hash_alg,
370        } => {
371            let storage: Option<&'static dyn StorageBackend> = match storage_type.as_str() {
372                "database" => {
373                    let db_storage = Box::new(DatabaseStorage::new(*storage_url.clone())?);
374                    Some(Box::leak(db_storage))
375                }
376                "rekor" => {
377                    let rekor_storage = Box::new(build_rekor_storage(
378                        *storage_url.clone(),
379                        &key,
380                        &cert,
381                        fulcio,
382                        &oidc_token,
383                    )?);
384                    Some(Box::leak(rekor_storage))
385                }
386                "local-fs" => {
387                    let fs_storage = Box::new(FilesystemStorage::new(storage_url.as_str())?);
388                    Some(Box::leak(fs_storage))
389                }
390                _ => None,
391            };
392
393            let config = ManifestCreationConfig {
394                paths: vec![path],
395                ingredient_names: vec!["Evaluation Results".to_string()],
396                name,
397                author_org,
398                author_name,
399                description,
400                linked_manifests: None,
401                storage,
402                print,
403                output_encoding: encoding,
404                key_path: key,
405                hash_alg: hash_alg.to_cose_algorithm(),
406                with_cc: false,
407                software_type: None,
408                version: None,
409                custom_fields: None,
410            };
411
412            manifest::evaluation::create_manifest(config, model_id, dataset_id, metrics)
413        }
414        EvaluationCommands::List {
415            storage_type,
416            storage_url,
417        } => {
418            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
419                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
420                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
421                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
422                _ => return Err(Error::Validation("Invalid storage type".to_string())),
423            };
424
425            manifest::evaluation::list_evaluation_manifests(storage.as_ref())
426        }
427        EvaluationCommands::Verify {
428            id,
429            storage_type,
430            storage_url,
431        } => {
432            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
433                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
434                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
435                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
436                _ => return Err(Error::Validation("Invalid storage type".to_string())),
437            };
438
439            manifest::evaluation::verify_evaluation_manifest(&id, storage.as_ref())
440        }
441    }
442}
443
444pub fn handle_cc_attestation_command(cmd: CCAttestationCommands) -> Result<()> {
445    match cmd {
446        CCAttestationCommands::Show => {
447            let _r = cc_attestation::get_report(true).unwrap();
448            Ok(())
449        }
450
451        CCAttestationCommands::GetLaunchMeasurement => {
452            let m = cc_attestation::get_launch_measurement().unwrap();
453            println!("Launch measurement raw bytes: 0x{}", hex::encode(m));
454            Ok(())
455        }
456
457        CCAttestationCommands::VerifyLaunch { host_platform } => {
458            let result = cc_attestation::verify_launch_endorsement(&host_platform).unwrap();
459            if result {
460                println!(
461                    "Passed: launch endorsement verification for {host_platform} host platform"
462                );
463            } else {
464                println!(
465                    "Failed: launch endorsement verification for {host_platform} host platform"
466                );
467            }
468            Ok(())
469        }
470    }
471}
472
473pub fn handle_software_command(cmd: SoftwareCommands) -> Result<()> {
474    match cmd {
475        SoftwareCommands::Create {
476            paths,
477            ingredient_names,
478            name,
479            software_type,
480            version,
481            author_org,
482            author_name,
483            description,
484            linked_manifests,
485            storage_type,
486            storage_url,
487            print,
488            encoding,
489            key,
490            cert,
491            fulcio,
492            oidc_token,
493            hash_alg,
494            with_tdx,
495        } => {
496            let storage: Option<&'static dyn StorageBackend> = match storage_type.as_str() {
497                "database" => {
498                    let db_storage = Box::new(DatabaseStorage::new(*storage_url.clone())?);
499                    Some(Box::leak(db_storage))
500                }
501                "rekor" => {
502                    let rekor_storage = Box::new(build_rekor_storage(
503                        *storage_url.clone(),
504                        &key,
505                        &cert,
506                        fulcio,
507                        &oidc_token,
508                    )?);
509                    Some(Box::leak(rekor_storage))
510                }
511                "local-fs" => {
512                    let fs_storage = Box::new(FilesystemStorage::new(storage_url.as_str())?);
513                    Some(Box::leak(fs_storage))
514                }
515                _ => None,
516            };
517
518            let config = ManifestCreationConfig {
519                paths,
520                ingredient_names,
521                name,
522                author_org,
523                author_name,
524                description,
525                linked_manifests,
526                storage,
527                print,
528                output_encoding: encoding,
529                key_path: key,
530                hash_alg: hash_alg.to_cose_algorithm(),
531                with_cc: with_tdx,
532                software_type: Some(software_type.clone()),
533                version: version.clone(),
534                custom_fields: None,
535            };
536
537            manifest::software::create_manifest(config, software_type, version)
538        }
539        SoftwareCommands::List {
540            storage_type,
541            storage_url,
542        } => {
543            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
544                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
545                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
546                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
547                _ => return Err(Error::Validation("Invalid storage type".to_string())),
548            };
549
550            manifest::software::list_software_manifests(storage.as_ref())
551        }
552        SoftwareCommands::Verify {
553            id,
554            storage_type,
555            storage_url,
556        } => {
557            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
558                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
559                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
560                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
561                _ => return Err(Error::Validation("Invalid storage type".to_string())),
562            };
563
564            manifest::software::verify_software_manifest(&id, storage.as_ref())
565        }
566        SoftwareCommands::LinkModel {
567            software_id,
568            model_id,
569            storage_type,
570            storage_url,
571        } => {
572            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
573                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
574                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
575                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
576                _ => return Err(Error::Validation("Invalid storage type".to_string())),
577            };
578
579            // Link software to model
580            manifest::link_manifests(&model_id, &software_id, storage.as_ref())
581        }
582        SoftwareCommands::LinkDataset {
583            software_id,
584            dataset_id,
585            storage_type,
586            storage_url,
587        } => {
588            let storage: Box<dyn StorageBackend> = match storage_type.as_str() {
589                "database" => Box::new(DatabaseStorage::new(*storage_url.clone())?),
590                "rekor" => Box::new(RekorStorage::new_with_url(*storage_url.clone())?),
591                "local-fs" => Box::new(FilesystemStorage::new(storage_url.as_str())?),
592                _ => return Err(Error::Validation("Invalid storage type".to_string())),
593            };
594
595            // Link software to dataset
596            manifest::link_manifests(&dataset_id, &software_id, storage.as_ref())
597        }
598    }
599}
600
601pub fn handle_rekor_command(cmd: RekorCommands) -> Result<()> {
602    match cmd {
603        RekorCommands::Verify {
604            uuid,
605            manifest,
606            rekor_url,
607        } => {
608            let manifest_bytes = std::fs::read(&manifest).map_err(|e| {
609                Error::Storage(format!(
610                    "Failed to read manifest file {}: {e}",
611                    manifest.display()
612                ))
613            })?;
614
615            let storage = RekorStorage::new_with_url(rekor_url)?;
616            let result = storage.verify_manifest(&manifest_bytes, &uuid)?;
617
618            println!("Rekor Verification Result:");
619            println!("  Log index:       {}", result.log_index);
620            println!("  Integrated time: {}", result.integrated_time);
621            if let Some(identity) = &result.signer_identity {
622                println!("  Signer identity: {identity}");
623            }
624            println!(
625                "  Payload hash:    {}",
626                if result.payload_hash_match {
627                    "MATCH"
628                } else {
629                    "MISMATCH"
630                }
631            );
632            println!(
633                "  Signature:       {}",
634                if result.signature_valid {
635                    "VALID"
636                } else {
637                    "INVALID"
638                }
639            );
640
641            if result.payload_hash_match && result.signature_valid {
642                println!("\nVerification PASSED: manifest matches Rekor entry.");
643                Ok(())
644            } else {
645                Err(Error::Validation(
646                    "Verification FAILED: manifest does not match Rekor entry.".to_string(),
647                ))
648            }
649        }
650        RekorCommands::Get { uuid, rekor_url } => {
651            let storage = RekorStorage::new_with_url(rekor_url)?;
652            let entry = storage.get_rekor_entry(&uuid)?;
653
654            println!("Rekor Entry:");
655            let uuid_str = entry.uuid.to_string();
656            let masked_uuid = if uuid_str.len() > 12 {
657                format!("{}...{}", &uuid_str[..8], &uuid_str[uuid_str.len() - 4..])
658            } else {
659                "[REDACTED]".to_string()
660            };
661            println!("  UUID:            {}", masked_uuid);
662            println!("  Log index:       {}", entry.log_index);
663            println!("  Integrated time: {}", entry.integrated_time);
664            println!("  Log ID:          {}", entry.log_id);
665
666            Ok(())
667        }
668    }
669}
670
671pub fn handle_pipeline_command(cmd: PipelineCommands) -> Result<()> {
672    match cmd {
673        PipelineCommands::GenerateProvenance {
674            inputs,
675            pipeline,
676            products,
677            key,
678            hash_alg,
679            encoding,
680            print,
681            storage_type,
682            storage_url,
683            with_tdx,
684        } => {
685            let storage: Option<&'static dyn StorageBackend> = match storage_type.as_str() {
686                "local-fs" => {
687                    let fs_storage = Box::new(FilesystemStorage::new(storage_url.as_str())?);
688                    Some(Box::leak(fs_storage))
689                }
690                _ => None,
691            };
692
693            slsa::cli::generate_build_provenance(
694                inputs,
695                pipeline,
696                products,
697                key,
698                hash_alg.to_cose_algorithm(),
699                encoding,
700                print,
701                storage,
702                with_tdx,
703            )
704        }
705    }
706}