icecave 0.2.0

use std::io::{Read, Write};
use std::os::fd::{FromRawFd, RawFd};
use std::sync::OnceLock;

// Minimal AF_VSOCK bindings and TCP fallback
const AF_VSOCK: i32 = 40; // Linux AF_VSOCK
const SOCK_STREAM: i32 = 1; // SOCK_STREAM
const VMADDR_CID_HOST: u32 = 3; // Host CID in Nitro Enclaves
const SHUT_WR: i32 = 1; // for shutdown(fd, SHUT_WR)

#[repr(C)]
struct SockAddrVm { svm_family: u16, svm_reserved1: u16, svm_port: u32, svm_cid: u32, svm_zero: [u8; 4] }
#[repr(C)]
struct SockAddr { sa_family: u16, sa_data: [u8; 14] }

extern "C" {
    fn socket(domain: i32, ty: i32, protocol: i32) -> i32;
    fn connect(sockfd: i32, addr: *const SockAddr, addrlen: u32) -> i32;
    fn shutdown(sockfd: i32, how: i32) -> i32;
}

static IDENTITY_SK: OnceLock<k256::SecretKey> = OnceLock::new();
static PCR0_CACHE: OnceLock<Vec<u8>> = OnceLock::new();
static PCR1_CACHE: OnceLock<Vec<u8>> = OnceLock::new();
static PCR2_CACHE: OnceLock<Vec<u8>> = OnceLock::new();

const DEFAULT_REPLICA: &str = "https://ic0.app";
const ATTESTATION_CANISTER_DEFAULT: &str = "ic6ec-6yaaa-aaaal-qss6q-cai";

fn gate_canister_id() -> String {
    std::env::var("ATTESTATION_GATE").unwrap_or_else(|_| ATTESTATION_CANISTER_DEFAULT.to_string())
}

pub struct VsockTransport { base: String }

impl VsockTransport {
    pub fn new(base: &str) -> Self { Self { base: base.to_string() } }

    fn forward(&self, method: &str, endpoint: &str, body: Option<&[u8]>) -> Result<(u16, Vec<u8>), String> {
        use serde_json::json;
        use std::env;
        let url = format!("{}/api/v2/{}", self.base.trim_end_matches('/'), endpoint.trim_start_matches('/'));
        let body_b64 = body.map(|b| base64::encode(b));
        let req = json!({
            "op": "forward",
            "url": url,
            "method": method,
            "headers": [],
            "body_b64": body_b64,
        }).to_string();

        if let Ok(addr) = env::var("PARENT_ADDR") {
            let mut stream = std::net::TcpStream::connect(addr).map_err(|e| e.to_string())?;
            stream.write_all(req.as_bytes()).map_err(|e| e.to_string())?;
            stream.flush().map_err(|e| e.to_string())?;
            stream.shutdown(std::net::Shutdown::Write).map_err(|e| e.to_string())?;
            let mut resp = Vec::with_capacity(4096);
            let mut buf = [0u8;4096];
            loop { match stream.read(&mut buf) { Ok(0)=>break, Ok(n)=>resp.extend_from_slice(&buf[..n]), Err(e)=>return Err(e.to_string()) } }
            let v: serde_json::Value = serde_json::from_slice(&resp).map_err(|e| e.to_string())?;
            if let Some(err) = v.get("error").and_then(|e| e.as_str()) { return Err(err.to_string()); }
            let status = v.get("status").and_then(|x| x.as_u64()).unwrap_or(0) as u16;
            let body_b64 = v.get("body_b64").and_then(|x| x.as_str()).unwrap_or("");
            let body = base64::decode(body_b64).map_err(|e| e.to_string())?;
            return Ok((status, body));
        }

        let fd = unsafe { socket(AF_VSOCK, SOCK_STREAM, 0) };
        if fd < 0 { return Err("socket() failed".to_string()); }
        let addr = SockAddrVm { svm_family: AF_VSOCK as u16, svm_reserved1: 0, svm_port: 7000, svm_cid: VMADDR_CID_HOST, svm_zero: [0;4] };
        let rc = unsafe { connect(fd, &addr as *const SockAddrVm as *const SockAddr, std::mem::size_of::<SockAddrVm>() as u32) };
        if rc != 0 { return Err("connect() failed".to_string()); }
        let mut file = unsafe { std::fs::File::from_raw_fd(fd as RawFd) };
        file.write_all(req.as_bytes()).map_err(|e| e.to_string())?;
        file.flush().map_err(|e| e.to_string())?;
        unsafe { let _ = shutdown(fd, SHUT_WR); }
        let mut resp = Vec::with_capacity(4096);
        let mut buf = [0u8;4096];
        loop { match file.read(&mut buf) { Ok(0)=>break, Ok(n)=>resp.extend_from_slice(&buf[..n]), Err(e)=>return Err(e.to_string()) } }
        let v: serde_json::Value = serde_json::from_slice(&resp).map_err(|e| e.to_string())?;
        if let Some(err) = v.get("error").and_then(|e| e.as_str()) { return Err(err.to_string()); }
        let status = v.get("status").and_then(|x| x.as_u64()).unwrap_or(0) as u16;
        let body_b64 = v.get("body_b64").and_then(|x| x.as_str()).unwrap_or("");
        let body = base64::decode(body_b64).map_err(|e| e.to_string())?;
        Ok((status, body))
    }
}

impl ic_agent::agent::Transport for VsockTransport {
    fn call(&self, effective_canister_id: ic_agent::export::Principal, envelope: Vec<u8>, _request_id: ic_agent::RequestId) -> std::pin::Pin<Box<dyn core::future::Future<Output = Result<(), ic_agent::AgentError>> + Send>> {
        let base = self.base.clone();
        Box::pin(async move {
            let me = VsockTransport { base };
            let endpoint = format!("canister/{}/call", effective_canister_id);
            let (status, _body) = me.forward("POST", &endpoint, Some(&envelope)).map_err(|s| ic_agent::AgentError::TransportError(Box::new(std::io::Error::new(std::io::ErrorKind::Other, s))))?;
            if status == 202 || status == 200 { Ok(()) } else { Err(ic_agent::AgentError::HttpError(ic_agent::agent::agent_error::HttpErrorPayload{ status: status as u16, content_type: None, content: Vec::new()})) }
        })
    }
    fn read_state(&self, effective_canister_id: ic_agent::export::Principal, envelope: Vec<u8>) -> std::pin::Pin<Box<dyn core::future::Future<Output = Result<Vec<u8>, ic_agent::AgentError>> + Send>> {
        let base = self.base.clone();
        Box::pin(async move {
            let me = VsockTransport { base };
            let endpoint = format!("canister/{}/read_state", effective_canister_id);
            let (_status, body) = me.forward("POST", &endpoint, Some(&envelope)).map_err(|s| ic_agent::AgentError::TransportError(Box::new(std::io::Error::new(std::io::ErrorKind::Other, s))))?;
            Ok(body)
        })
    }
    fn read_subnet_state(&self, subnet_id: ic_agent::export::Principal, envelope: Vec<u8>) -> std::pin::Pin<Box<dyn core::future::Future<Output = Result<Vec<u8>, ic_agent::AgentError>> + Send>> {
        let base = self.base.clone();
        Box::pin(async move {
            let me = VsockTransport { base };
            let endpoint = format!("subnet/{}/read_state", subnet_id);
            let (_status, body) = me.forward("POST", &endpoint, Some(&envelope)).map_err(|s| ic_agent::AgentError::TransportError(Box::new(std::io::Error::new(std::io::ErrorKind::Other, s))))?;
            Ok(body)
        })
    }
    fn query(&self, effective_canister_id: ic_agent::export::Principal, envelope: Vec<u8>) -> std::pin::Pin<Box<dyn core::future::Future<Output = Result<Vec<u8>, ic_agent::AgentError>> + Send>> {
        let base = self.base.clone();
        Box::pin(async move {
            let me = VsockTransport { base };
            let endpoint = format!("canister/{}/query", effective_canister_id);
            let (_status, body) = me.forward("POST", &endpoint, Some(&envelope)).map_err(|s| ic_agent::AgentError::TransportError(Box::new(std::io::Error::new(std::io::ErrorKind::Other, s))))?;
            Ok(body)
        })
    }
    fn status(&self) -> std::pin::Pin<Box<dyn core::future::Future<Output = Result<Vec<u8>, ic_agent::AgentError>> + Send>> {
        let base = self.base.clone();
        Box::pin(async move {
            let me = VsockTransport { base };
            let (_status, body) = me.forward("GET", "status", None).map_err(|s| ic_agent::AgentError::TransportError(Box::new(std::io::Error::new(std::io::ErrorKind::Other, s))))?;
            Ok(body)
        })
    }
}

fn shared_identity() -> ic_agent::identity::Secp256k1Identity {
    use ic_agent::identity::Secp256k1Identity;
    let sk = IDENTITY_SK.get_or_init(|| k256::SecretKey::random(&mut rand::thread_rng()));
    let bytes = sk.to_bytes();
    let k = k256::SecretKey::from_bytes(&bytes).expect("secp key");
    Secp256k1Identity::from_private_key(k)
}

fn build_agent_with(id: &ic_agent::identity::Secp256k1Identity) -> Result<ic_agent::Agent, String> {
    let transport = VsockTransport::new(DEFAULT_REPLICA);
    ic_agent::Agent::builder().with_transport(transport).with_identity(id.clone()).build().map_err(|e| e.to_string())
}

fn check_ic_status(agent: &ic_agent::Agent) -> Result<(), String> {
    let rt = tokio::runtime::Builder::new_current_thread().enable_time().build().map_err(|e| e.to_string())?;
    let _ = rt.block_on(async { agent.status().await }).map_err(|e| e.to_string())?;
    Ok(())
}

fn get_challenge_via_agent(agent: &ic_agent::Agent, canister: &str, project: &str) -> Result<Vec<u8>, String> {
    use candid::{Encode, Decode};
    let cid: candid::Principal = canister.parse::<candid::Principal>().map_err(|e| e.to_string())?;
    let arg = Encode!(&project.to_string()).map_err(|e| e.to_string())?;
    let rt = tokio::runtime::Builder::new_current_thread().enable_time().build().map_err(|e| e.to_string())?;
    eprintln!("[icecave] requesting attestation challenge…");
    let out = rt.block_on(async { agent.update(&cid, "request_attestation_challenge").with_arg(arg).call_and_wait().await }).map_err(|e| {
        eprintln!("[icecave] challenge request error: {}", e);
        e.to_string()
    })?;
    #[derive(candid::CandidType, serde::Deserialize)]
    struct Challenge { bytes: Vec<u8>, issued_at_ns: u64 }
    let ch: Challenge = Decode!(&out, Challenge).map_err(|e| e.to_string())?;
    eprintln!("[icecave] challenge received: len={} issued_at_ns={}", ch.bytes.len(), ch.issued_at_ns);
    Ok(ch.bytes)
}

fn nitro_attest_with_user_data(user_data: Vec<u8>) -> Result<Vec<u8>, String> {
    use aws_nitro_enclaves_nsm_api::{api::Request, driver};
    let fd = driver::nsm_init();
    if fd < 0 { return Err("nsm_init failed".into()); }
    eprintln!("[icecave] NSM attestation: starting (user_data_len={})", user_data.len());
    let req = Request::Attestation { user_data: Some(user_data.into()), nonce: None, public_key: None };
    let resp = driver::nsm_process_request(fd, req);
    driver::nsm_exit(fd);
    match resp {
        aws_nitro_enclaves_nsm_api::api::Response::Attestation { document } => {
            eprintln!("[icecave] NSM attestation: document len={}", document.len());
            // Best-effort: log PCR0..2 from the attestation document's payload for comparison.
            log_pcrs_from_doc(&document);
            log_user_data_from_doc(&document);
            Ok(document)
        },
        aws_nitro_enclaves_nsm_api::api::Response::Error(e) => Err(format!("nsm attestation error: {:?}", e)),
        _ => Err("unexpected nsm response".to_string()),
    }
}

fn log_user_data_from_doc(doc: &[u8]) {
    // Parse COSE_Sign1 and payload, then extract and print user_data summary
    let top: serde_cbor::Value = match serde_cbor::from_slice(doc) { Ok(v) => v, Err(e) => { eprintln!("[icecave] user_data: COSE decode error: {}", e); return; } };
    let payload_b = match top {
        serde_cbor::Value::Array(ref arr) if arr.len() == 4 => match arr.get(2) { Some(serde_cbor::Value::Bytes(b)) => b.clone(), _ => { eprintln!("[icecave] user_data: no payload bstr"); return; }},
        _ => { eprintln!("[icecave] user_data: not COSE Sign1 array"); return; }
    };
    let inner: serde_cbor::Value = match serde_cbor::from_slice(&payload_b) { Ok(v) => v, Err(e) => { eprintln!("[icecave] user_data: payload cbor error: {}", e); return; } };
    let m = match inner { serde_cbor::Value::Map(m) => m, _ => { eprintln!("[icecave] user_data: payload not a map"); return; } };
    for (k, v) in &m { if let serde_cbor::Value::Text(t) = k { if t == "user_data" {
            match v {
                serde_cbor::Value::Bytes(b) => {
                    eprintln!("[icecave] user_data: len={}", b.len());
                    // Try to decode as CBOR map and print keys
                    if let Ok(val) = serde_cbor::from_slice::<serde_cbor::Value>(b) {
                        match val {
                            serde_cbor::Value::Map(mm) => {
                                let mut keys: Vec<String> = Vec::new();
                                for (kk, vv) in mm { let key = match kk { serde_cbor::Value::Text(s)=>s, _=>"<non-text>".to_string() }; let kind = match vv { serde_cbor::Value::Bytes(_)=>"bytes", serde_cbor::Value::Text(_)=>"text", serde_cbor::Value::Array(_)=>"array", _=>"other" }; keys.push(format!("{}:{}", key, kind)); }
                                eprintln!("[icecave] user_data: map keys=[{}]", keys.join(", "));
                            }
                            other => { eprintln!("[icecave] user_data: not map, kind={:?}", other); }
                        }
                    } else {
                        eprintln!("[icecave] user_data: not CBOR");
                    }
                }
                other => { eprintln!("[icecave] user_data: not bytes kind={:?}", other); }
            }
        } } }
}

// projectless attestation removed

fn submit_attestation_via_agent_project(agent: &ic_agent::Agent, canister: &str, project: &str, doc: &[u8], challenge: &[u8]) -> Result<(), String> {
    use candid::{Encode, Decode};
    let cid: candid::Principal = canister.parse::<candid::Principal>().map_err(|e| e.to_string())?;
    let arg = Encode!(&project.to_string(), &doc.to_vec(), &challenge.to_vec()).map_err(|e| e.to_string())?;
    let rt = tokio::runtime::Builder::new_current_thread().enable_time().build().map_err(|e| e.to_string())?;
    eprintln!("[icecave] submitting attestation for project='{}' (doc_len={}, chal_len={})", project, doc.len(), challenge.len());
    let out = rt.block_on(async { agent.update(&cid, "submit_nitro_attestation_for_project").with_arg(arg).call_and_wait().await }).map_err(|e| {
        eprintln!("[icecave] submit error: {}", e);
        e.to_string()
    })?;
    let result: Result<bool, String> = Decode!(&out, Result<bool, String>).map_err(|e| {
        eprintln!("[icecave] decode error: {}", e);
        e.to_string()
    })?;
    match result {
        Ok(true) => { eprintln!("[icecave] submit_ok (project={})", project); Ok(()) },
        Ok(false) => { eprintln!("[icecave] submit returned false (project={})", project); Err("submit_for_project returned false".into()) },
        Err(e) => { eprintln!("[icecave] submit error (project={}): {}", project, e); Err(e) },
    }
}

pub struct IcecaveReady {
    identity: ic_agent::identity::Secp256k1Identity,
}

impl IcecaveReady {
    pub fn identity(&self) -> ic_agent::identity::Secp256k1Identity { self.identity.clone() }
    pub fn agent(&self) -> Result<ic_agent::Agent, String> { build_agent_with(&self.identity) }
}

/// Start Icecave: generate identity, check IC status, attempt attestation up to 10 times (1 minute between).
/// On success, invokes `on_ready` with identity and returns Ok(()) after the callback completes.
// projectless start removed

/// Start Icecave attestation for a specific Gate project.
pub fn start_with_project<F>(project: &str, on_ready: F) -> Result<(), String>
where
    F: FnOnce(IcecaveReady),
{
    let id = shared_identity();
    let project = project.to_string();
    let mut last_err: Option<String> = None;
    for attempt in 1..=10u32 {
        eprintln!("[icecave] attempt {}: checking IC status…", attempt);
        let agent = match build_agent_with(&id) {
            Ok(a) => a,
            Err(e) => { eprintln!("[icecave] build agent failed: {}", e); last_err = Some(e); backoff_if_needed(attempt)?; continue; }
        };
        if let Err(e) = check_ic_status(&agent) {
            eprintln!("[icecave] status check failed: {}", e);
            last_err = Some(e);
            backoff_if_needed(attempt)?;
            continue;
        }
        eprintln!("[icecave] attempt {}: IC reachable, attesting (project={})…", attempt, project);
        match do_attest_with_project(&agent, &project) {
            Ok(()) => {
                eprintln!("[icecave] attestation success (project={})", project);
                on_ready(IcecaveReady { identity: id.clone() });
                return Ok(());
            }
            Err(e) => {
                eprintln!("[icecave] attestation failed: {}", e);
                last_err = Some(e);
                backoff_if_needed(attempt)?;
            }
        }
    }
    Err(last_err.unwrap_or_else(|| "attestation failed".to_string()))
}

fn backoff_if_needed(attempt: u32) -> Result<(), String> {
    if attempt < 10 {
        let rt = tokio::runtime::Builder::new_current_thread().enable_time().build().map_err(|e| e.to_string())?;
        rt.block_on(async { tokio::time::sleep(std::time::Duration::from_secs(60)).await; });
    }
    Ok(())
}

// projectless attestation removed

fn do_attest_with_project(agent: &ic_agent::Agent, project: &str) -> Result<(), String> {
    use ic_agent::Identity;
    let gate = gate_canister_id();
    let challenge = get_challenge_via_agent(agent, &gate, project)?;
    let id = shared_identity();
    let principal_text = id.sender().map_err(|e| e.to_string())?.to_text();
    let user_data = encode_user_data(&challenge, &principal_text)?;
    let doc = nitro_attest_with_user_data(user_data)?;
    submit_attestation_via_agent_project(agent, &gate, project, &doc, &challenge)?;
    // On success, persist principal to host for CLI discovery
    let _ = save_principal_to_parent(project, &principal_text);
    Ok(())
}

// Public helpers
pub fn default_replica() -> &'static str { DEFAULT_REPLICA }
pub fn vsock_transport() -> VsockTransport { VsockTransport::new(DEFAULT_REPLICA) }

fn hex_bytes(bs: &[u8]) -> String {
    const HEX: &[u8; 16] = b"0123456789abcdef";
    let mut out = String::with_capacity(bs.len()*2);
    for &b in bs { out.push(HEX[(b>>4) as usize] as char); out.push(HEX[(b&0xF) as usize] as char); }
    out
}

fn log_pcrs_from_doc(doc: &[u8]) {
    // COSE_Sign1 structure: [ protected, unprotected, payload_bstr, signature ]
    let top: serde_cbor::Value = match serde_cbor::from_slice(doc) { Ok(v) => v, Err(_) => return };
    let payload_b = match top { serde_cbor::Value::Array(ref arr) if arr.len() == 4 => {
            match arr.get(2) { Some(serde_cbor::Value::Bytes(b)) => b.clone(), _ => return }
        }, _ => return };
    let inner: serde_cbor::Value = match serde_cbor::from_slice(&payload_b) { Ok(v) => v, Err(_) => return };
    let pcrs_map = match inner {
        serde_cbor::Value::Map(map) => {
            // Find key "pcrs" in the map
            let mut found: Option<serde_cbor::value::Value> = None;
            for (k,v) in &map {
                if let serde_cbor::Value::Text(s) = k { if s == "pcrs" { found = Some(v.clone()); break; } }
            }
            found
        },
        _ => None,
    };
    let m = match pcrs_map { Some(serde_cbor::Value::Map(m)) => m, _ => return };
    let get_bytes = |key: i128| -> Option<Vec<u8>> {
        for (k, v) in &m {
            if let serde_cbor::Value::Integer(i) = k {
                if *i == key {
                    match v {
                        serde_cbor::Value::Bytes(b) => return Some(b.clone()),
                        serde_cbor::Value::Array(arr) => {
                            if let Some(serde_cbor::Value::Bytes(b)) = arr.first() { return Some(b.clone()); }
                        }
                        _ => {}
                    }
                }
            }
        }
        None
    };
    if let Some(p0) = get_bytes(0) { let _ = PCR0_CACHE.set(p0.clone()); eprintln!("[icecave] PCR0={}", hex_bytes(&p0)); }
    if let Some(p1) = get_bytes(1) { let _ = PCR1_CACHE.set(p1.clone()); eprintln!("[icecave] PCR1={}", hex_bytes(&p1)); }
    if let Some(p2) = get_bytes(2) { let _ = PCR2_CACHE.set(p2.clone()); eprintln!("[icecave] PCR2={}", hex_bytes(&p2)); }
}

fn encode_user_data(challenge: &[u8], principal_text: &str) -> Result<Vec<u8>, String> {
    use serde::Serialize;
    #[derive(Serialize)]
    struct Ud<'a> {
        #[serde(with = "serde_bytes")]
        challenge: &'a [u8],
        principal: &'a str,
    }
    serde_cbor::to_vec(&Ud { challenge, principal: principal_text }).map_err(|e| e.to_string())
}

fn save_principal_to_parent(project: &str, principal: &str) -> Result<(), String> {
    use serde_json::json;
    use std::env;
    let req = json!({ "op": "save_principal", "project": project, "principal": principal }).to_string();

    if let Ok(addr) = env::var("PARENT_ADDR") {
        let mut stream = std::net::TcpStream::connect(addr).map_err(|e| e.to_string())?;
        stream.write_all(req.as_bytes()).map_err(|e| e.to_string())?;
        stream.flush().map_err(|e| e.to_string())?;
        stream.shutdown(std::net::Shutdown::Write).map_err(|e| e.to_string())?;
        let mut buf = [0u8; 1024];
        let _ = stream.read(&mut buf).unwrap_or(0);
        return Ok(());
    }

    // vsock path
    let fd = unsafe { socket(AF_VSOCK, SOCK_STREAM, 0) };
    if fd < 0 { return Err("socket() failed".to_string()); }
    let addr = SockAddrVm { svm_family: AF_VSOCK as u16, svm_reserved1: 0, svm_port: 7000, svm_cid: VMADDR_CID_HOST, svm_zero: [0;4] };
    let rc = unsafe { connect(fd, &addr as *const SockAddrVm as *const SockAddr, std::mem::size_of::<SockAddrVm>() as u32) };
    if rc != 0 { return Err("connect() failed".to_string()); }
    let mut file = unsafe { std::fs::File::from_raw_fd(fd as RawFd) };
    file.write_all(req.as_bytes()).map_err(|e| e.to_string())?;
    file.flush().map_err(|e| e.to_string())?;
    unsafe { let _ = shutdown(fd, SHUT_WR); }
    let mut buf = [0u8;1024];
    let _ = file.read(&mut buf).unwrap_or(0);
    Ok(())
}

/// Returns cached PCR values (0,1,2,15) from the last attestation document, as lowercase hex.
pub fn pcr0_hex() -> Option<String> { PCR0_CACHE.get().map(|b| hex_bytes(b)) }
pub fn pcr1_hex() -> Option<String> { PCR1_CACHE.get().map(|b| hex_bytes(b)) }
pub fn pcr2_hex() -> Option<String> { PCR2_CACHE.get().map(|b| hex_bytes(b)) }