Skip to content

Commit ce6d65e

Browse files
sarroutbisarroutbi
committed
[DRAFT]: Code refactoring
Signed-off-by: Sergio Arroutbi <[email protected]>
1 parent be33d8b commit ce6d65e

File tree

1 file changed

+126
-91
lines changed

1 file changed

+126
-91
lines changed

‎keylime-agent/src/main.rs

+126-91
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,20 @@ pub struct QuoteData<'a> {
125125
work_dir: PathBuf,
126126
}
127127

128+
#[derive(Debug)]
129+
pub struct AgentActivation {
130+
ak: tpm::AKResult,
131+
ek_result: tpm::EKResult,
132+
api_versions: Vec<String>,
133+
agent: config::AgentConfig,
134+
agent_uuid: String,
135+
mtls_cert: Option<X509>,
136+
device_id: Option<keylime::device_id::DeviceID>,
137+
attest: Option<tss_esapi::structures::Attest>,
138+
signature: Option<tss_esapi::structures::Signature>,
139+
ak_handle: KeyHandle,
140+
}
141+
128142
#[actix_web::main]
129143
async fn main() -> Result<()> {
130144
// Print --help information
@@ -611,98 +625,23 @@ async fn main() -> Result<()> {
611625
warn!("mTLS disabled, Tenant and Verifier will reach out to agent via HTTP");
612626
}
613627

614-
{
615-
// Declare here as these must live longer than the builder
616-
let iak_pub;
617-
let idevid_pub;
618-
let ak_pub = &PublicBuffer::try_from(ak.public)?.marshall()?;
619-
let ek_pub =
620-
&PublicBuffer::try_from(ek_result.public.clone())?.marshall()?;
621-
622-
// Create a RegistrarClientBuilder and set the parameters
623-
let mut builder = RegistrarClientBuilder::new()
624-
.ak_pub(ak_pub)
625-
.ek_pub(ek_pub)
626-
.enabled_api_versions(
627-
api_versions.iter().map(|ver| ver.as_ref()).collect(),
628-
)
629-
.registrar_ip(config.agent.registrar_ip.clone())
630-
.registrar_port(config.agent.registrar_port)
631-
.uuid(&agent_uuid)
632-
.ip(config.agent.contact_ip.clone())
633-
.port(config.agent.contact_port);
634-
635-
if let Some(mtls_cert) = mtls_cert {
636-
builder = builder.mtls_cert(mtls_cert);
637-
}
638-
639-
// If the certificate is not None add it to the builder
640-
if let Some(ek_cert) = ek_result.ek_cert {
641-
builder = builder.ek_cert(ek_cert);
642-
}
643-
644-
// Set the IAK/IDevID related fields, if enabled
645-
if config.agent.enable_iak_idevid {
646-
let (Some(dev_id), Some(attest), Some(signature)) =
647-
(&device_id, attest, signature)
648-
else {
649-
error!(
650-
"IDevID and IAK are enabled but could not be generated"
651-
);
652-
return Err(Error::Configuration(config::KeylimeConfigError::Generic(
653-
"IDevID and IAK are enabled but could not be generated"
654-
.to_string(),
655-
)));
656-
};
657-
658-
iak_pub = PublicBuffer::try_from(dev_id.iak_pubkey.clone())?
659-
.marshall()?;
660-
idevid_pub =
661-
PublicBuffer::try_from(dev_id.idevid_pubkey.clone())?
662-
.marshall()?;
663-
builder = builder
664-
.iak_attest(attest.marshall()?)
665-
.iak_sign(signature.marshall()?)
666-
.iak_pub(&iak_pub)
667-
.idevid_pub(&idevid_pub);
668-
669-
// If the IAK certificate was provided, set it
670-
if let Some(iak_cert) = dev_id.iak_cert.clone() {
671-
builder = builder.iak_cert(iak_cert);
672-
}
673-
674-
// If the IDevID certificate was provided, set it
675-
if let Some(idevid_cert) = dev_id.idevid_cert.clone() {
676-
builder = builder.idevid_cert(idevid_cert);
677-
}
678-
}
679-
680-
// Build the registrar client
681-
let mut registrar_client = builder.build().await?;
682-
683-
// Request keyblob material
684-
let keyblob = registrar_client.register_agent().await?;
685-
686-
info!("SUCCESS: Agent {} registered", &agent_uuid);
687-
688-
let key = ctx.activate_credential(
689-
keyblob,
690-
ak_handle,
691-
ek_result.key_handle,
692-
)?;
693-
694-
// Flush EK if we created it
695-
if config.agent.ek_handle.is_empty() {
696-
ctx.flush_context(ek_result.key_handle.into())?;
628+
let aa = AgentActivation {
629+
ak,
630+
ek_result,
631+
api_versions: api_versions.clone(),
632+
agent: config.agent.clone(),
633+
agent_uuid: agent_uuid.clone(),
634+
mtls_cert,
635+
device_id,
636+
attest,
637+
signature,
638+
ak_handle,
639+
};
640+
match activate_agent(aa, &mut ctx).await {
641+
Ok(()) => (),
642+
Err(e) => {
643+
error!("Failed to activate agent: {}", e);
697644
}
698-
699-
let mackey = general_purpose::STANDARD.encode(key.value());
700-
let auth_tag =
701-
crypto::compute_hmac(mackey.as_bytes(), agent_uuid.as_bytes())?;
702-
let auth_tag = hex::encode(&auth_tag);
703-
704-
registrar_client.activate_agent(&auth_tag).await?;
705-
info!("SUCCESS: Agent {} activated", &agent_uuid);
706645
}
707646

708647
let (mut payload_tx, mut payload_rx) =
@@ -945,6 +884,102 @@ async fn main() -> Result<()> {
945884
result.map(|_| ())
946885
}
947886

887+
async fn activate_agent(
888+
mut aa: AgentActivation,
889+
mut ctx: &mut tpm::Context<'_>,
890+
) -> Result<()> {
891+
let iak_pub;
892+
let idevid_pub;
893+
let ak_pub = &PublicBuffer::try_from(aa.ak.public)?.marshall()?;
894+
let ek_pub =
895+
&PublicBuffer::try_from(aa.ek_result.public.clone())?.marshall()?;
896+
897+
// Create a RegistrarClientBuilder and set the parameters
898+
let mut builder = RegistrarClientBuilder::new()
899+
.ak_pub(ak_pub)
900+
.ek_pub(ek_pub)
901+
.enabled_api_versions(
902+
aa.api_versions.iter().map(|ver| ver.as_ref()).collect(),
903+
)
904+
.registrar_ip(aa.agent.registrar_ip.clone())
905+
.registrar_port(aa.agent.registrar_port)
906+
.uuid(&aa.agent_uuid)
907+
.ip(aa.agent.contact_ip.clone())
908+
.port(aa.agent.contact_port);
909+
910+
if let Some(mtls_cert) = aa.mtls_cert {
911+
builder = builder.mtls_cert(mtls_cert);
912+
}
913+
914+
// If the certificate is not None add it to the builder
915+
if let Some(ek_cert) = aa.ek_result.ek_cert {
916+
builder = builder.ek_cert(ek_cert);
917+
}
918+
919+
// Set the IAK/IDevID related fields, if enabled
920+
if aa.agent.enable_iak_idevid {
921+
let (Some(dev_id), Some(attest), Some(signature)) =
922+
(&aa.device_id, aa.attest, aa.signature)
923+
else {
924+
error!("IDevID and IAK are enabled but could not be generated");
925+
return Err(Error::Configuration(
926+
config::KeylimeConfigError::Generic(
927+
"IDevID and IAK are enabled but could not be generated"
928+
.to_string(),
929+
),
930+
));
931+
};
932+
933+
iak_pub =
934+
PublicBuffer::try_from(dev_id.iak_pubkey.clone())?.marshall()?;
935+
idevid_pub = PublicBuffer::try_from(dev_id.idevid_pubkey.clone())?
936+
.marshall()?;
937+
builder = builder
938+
.iak_attest(attest.marshall()?)
939+
.iak_sign(signature.marshall()?)
940+
.iak_pub(&iak_pub)
941+
.idevid_pub(&idevid_pub);
942+
943+
// If the IAK certificate was provided, set it
944+
if let Some(iak_cert) = dev_id.iak_cert.clone() {
945+
builder = builder.iak_cert(iak_cert);
946+
}
947+
948+
// If the IDevID certificate was provided, set it
949+
if let Some(idevid_cert) = dev_id.idevid_cert.clone() {
950+
builder = builder.idevid_cert(idevid_cert);
951+
}
952+
}
953+
954+
// Build the registrar client
955+
let mut registrar_client = builder.build().await?;
956+
957+
// Request keyblob material
958+
let keyblob = registrar_client.register_agent().await?;
959+
960+
info!("SUCCESS: Agent {} registered", &aa.agent_uuid);
961+
962+
let key = ctx.activate_credential(
963+
keyblob,
964+
aa.ak_handle,
965+
aa.ek_result.key_handle,
966+
)?;
967+
968+
// Flush EK if we created it
969+
if aa.agent.ek_handle.is_empty() {
970+
ctx.flush_context(aa.ek_result.key_handle.into())?;
971+
}
972+
973+
let mackey = general_purpose::STANDARD.encode(key.value());
974+
let auth_tag =
975+
crypto::compute_hmac(mackey.as_bytes(), aa.agent_uuid.as_bytes())?;
976+
let auth_tag = hex::encode(&auth_tag);
977+
978+
registrar_client.activate_agent(&auth_tag).await?;
979+
info!("SUCCESS: Agent {} activated", &aa.agent_uuid);
980+
Ok(())
981+
}
982+
948983
/*
949984
* Input: file path
950985
* Output: file content

0 commit comments

Comments
 (0)