// Copyright 2020-2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use std::fs::OpenOptions;
use std::io::{Read, Write};
use std::process::Command;
use std::rc::Rc;
use std::time::{Duration, Instant};

use log::{debug, error, info};
use nix::unistd;
use serde::{Deserialize, Serialize};
use serde_json;

use super::{httpd, nginx, PostSyncAction};
use crate::imds;
use crate::util;
use crate::{config, defs, enclave, enclave::P11neEnclave};
use vtok_rpc::api::schema;

#[derive(Debug)]
pub enum Error {
    AcmDbFetchError(Option<i32>, String),
    AcmDbParseError(serde_json::Error),
    AddTokenError(schema::ApiError),
    AwsCliExecError(std::io::Error),
    BadGroup(String),
    BadUser(String),
    EnclaveError(enclave::Error),
    FileDbError(std::io::Error),
    FileDbParseError(serde_json::Error),
    NixError(nix::Error),
    ImdsError(imds::Error),
    PinGenerationError(std::io::Error),
    RefreshTokenError(schema::ApiError),
    SourceDbEmpty,
    UpdateTokenError(schema::ApiError),
    TargetIoError(std::io::Error),
    ManagedServiceUnknown,
}

#[derive(Debug, Deserialize, PartialEq, Serialize)]
struct AcmDb {
    certificate: String,
    #[serde(rename = "certificateChain")]
    certificate_chain: String,
    #[serde(rename = "encryptedPrivateKey")]
    encrypted_private_key: String,
    #[serde(rename = "encryptionMethod")]
    encryption_method: String,
}

type FileDb = Vec<schema::PrivateKey>;

enum DbSource {
    Acm {
        cert_arn: String,
        db: AcmDb,
        bucket: String,
    },
    File {
        db: FileDb,
    },
}

impl DbSource {
    pub fn new(source_config: config::Source) -> Result<Self, Error> {
        match source_config {
            config::Source::Acm {
                certificate_arn,
                bucket,
            } => {
                let bucket = bucket.unwrap_or(defs::DEFAULT_ACM_BUCKET.to_string());
                Ok(Self::Acm {
                    db: Self::fetch_acm_db(certificate_arn.as_str(), bucket.as_str())?,
                    cert_arn: certificate_arn,
                    bucket,
                })
            }
            config::Source::FileDb { path } => Ok(Self::File {
                db: Self::fetch_file_db(path.as_str())?,
            }),
        }
    }

    pub fn update(&mut self) -> Result<bool, Error> {
        match self {
            Self::Acm {
                ref mut db,
                cert_arn,
                bucket,
            } => {
                let new_db = Self::fetch_acm_db(cert_arn, bucket)?;
                let res = new_db != *db;
                *db = new_db;
                Ok(res)
            }
            Self::File { .. } => Ok(false),
        }
    }

    pub fn to_schema_keys(&self) -> Vec<schema::PrivateKey> {
        match self {
            Self::Acm { db, .. } => {
                let optcerts = self.cert_pem().map(|c| {
                    let mut cc = c.to_string();
                    self.cert_chain_pem().map(|ch| cc.push_str(ch));
                    cc
                });
                vec![schema::PrivateKey {
                    id: 1,
                    label: "acm-key".to_string(),
                    encrypted_pem_b64: db.encrypted_private_key.clone(),
                    cert_pem: optcerts,
                }]
            }
            Self::File { db, .. } => db.as_slice().to_vec(),
        }
    }

    pub fn cert_pem(&self) -> Option<&str> {
        match self {
            Self::Acm { db, .. } => Some(db.certificate.as_str()),
            Self::File { .. } => None,
        }
    }

    pub fn cert_chain_pem(&self) -> Option<&str> {
        match self {
            Self::Acm { db, .. } => Some(db.certificate_chain.as_str()),
            Self::File { .. } => None,
        }
    }

    fn fetch_acm_db(cert_arn: &str, bucket: &str) -> Result<AcmDb, Error> {
        debug!("Fetching cert for arn: {}", cert_arn);
        let s3_url = format!(
            "s3://{}-ec2-enclave-certificate-{}-{}/{}/{}",
            imds::partition().map_err(Error::ImdsError)?,
            imds::region().map_err(Error::ImdsError)?,
            bucket,
            imds::role_arn().map_err(Error::ImdsError)?,
            cert_arn
        );
        let output = Command::new("aws")
            .arg("s3")
            .args(&[
                "--region",
                imds::region().map_err(Error::ImdsError)?.as_str(),
            ])
            .args(&["cp", s3_url.as_str(), "-"])
            .output()
            .map_err(Error::AwsCliExecError)?;
        if !output.status.success() {
            return Err(Error::AcmDbFetchError(
                output.status.code(),
                String::from_utf8_lossy(output.stderr.as_slice()).to_string(),
            ));
        }
        serde_json::from_slice::<AcmDb>(output.stdout.as_slice()).map_err(Error::AcmDbParseError)
    }

    fn fetch_file_db(path: &str) -> Result<FileDb, Error> {
        OpenOptions::new()
            .read(true)
            .open(path)
            .and_then(|mut f| {
                let mut buf = Vec::new();
                f.read_to_end(&mut buf)?;
                Ok(buf)
            })
            .map_err(Error::FileDbError)
            .and_then(|buf| {
                serde_json::from_slice::<FileDb>(buf.as_slice()).map_err(Error::FileDbParseError)
            })
    }
}

#[derive(Debug, Clone, Copy)]
pub enum ManagedService {
    Nginx,
    Httpd,
}

impl ManagedService {
    pub fn from_str(service: &str) -> Result<Self, Error> {
        match service {
            defs::SERVICE_NGINX => Ok(ManagedService::Nginx),
            defs::SERVICE_HTTPD => Ok(ManagedService::Httpd),
            _ => Err(Error::ManagedServiceUnknown),
        }
    }
}

pub struct ManagedToken {
    pub label: String,
    pub pin: String,
    db: DbSource,
    target: Option<config::Target>,
    enclave: Rc<P11neEnclave>,
    refresh_interval: Duration,
    next_refresh: Instant,
    service: ManagedService,
}

impl ManagedToken {
    pub fn new(
        token_config: config::Token,
        enclave: Rc<P11neEnclave>,
        service: ManagedService,
    ) -> Result<Self, Error> {
        let pin = match token_config.pin {
            Some(pin) => pin,
            None => util::generate_pkcs11_pin().map_err(Error::PinGenerationError)?,
        };
        Ok(Self {
            label: token_config.label,
            pin,
            db: DbSource::new(token_config.source)?,
            target: token_config.target,
            enclave,
            refresh_interval: Duration::from_secs(
                token_config
                    .refresh_interval_secs
                    .unwrap_or(defs::DEFAULT_TOKEN_REFRESH_INTERVAL_SECS),
            ),
            next_refresh: Instant::now(),
            service,
        })
    }

    pub fn sync(&mut self) -> Result<Option<PostSyncAction>, Error> {
        let db_changed = self.db.update()?;
        let is_online = self
            .enclave
            .describe_token(self.label.clone(), self.pin.clone())
            .map_err(Error::EnclaveError)?
            .is_ok();

        match (is_online, db_changed, self.target.as_ref()) {
            (true, true, Some(_target)) => {
                info!(
                    "Managed service certificate changed. Updating token {}",
                    self.label.as_str()
                );

                debug!("Updating token {}", self.label.as_str());
                self.enclave
                    .update_token(
                        self.label.clone(),
                        self.pin.clone(),
                        self.to_schema_token()?,
                    )
                    .map_err(Error::EnclaveError)?
                    .map_err(Error::UpdateTokenError)?;

                debug!("Updating managed service configuration");
                self.satisfy_target().or_else(|e| {
                    error!(
                        "Unable to satisfy target for token {}: {:?}",
                        self.label.as_str(),
                        e
                    );
                    Ok(None)
                })
            }
            (true, true, None) => {
                info!(
                    "DB change detected for token {}. Updating",
                    self.label.as_str()
                );
                self.enclave
                    .update_token(
                        self.label.clone(),
                        self.pin.clone(),
                        self.to_schema_token()?,
                    )
                    .map_err(Error::EnclaveError)?
                    .map_err(Error::UpdateTokenError)?;
                Ok(None)
            }
            (true, false, maybe_target) => {
                if self.next_refresh <= Instant::now() {
                    info!("Refreshing token {}", self.label.as_str());
                    self.enclave
                        .refresh_token(
                            self.label.clone(),
                            self.pin.clone(),
                            Self::kms_envelope_key()?,
                        )
                        .map_err(Error::EnclaveError)?
                        .map_err(Error::RefreshTokenError)?;
                    self.next_refresh += self.refresh_interval;
                    Ok(maybe_target.as_ref().map(|t| match t {
                        config::Target::Conf { .. } => match self.service {
                            ManagedService::Nginx => PostSyncAction::ReloadNginx,
                            ManagedService::Httpd => PostSyncAction::ReloadHttpd,
                        },
                    }))
                } else {
                    Ok(None)
                }
            }
            (false, _, _) => {
                debug!("Adding token {}", self.label.as_str());
                self.enclave
                    .add_token(self.to_schema_token()?)
                    .map_err(Error::EnclaveError)?
                    .map_err(Error::AddTokenError)?;
                self.satisfy_target().or_else(|e| {
                    error!(
                        "Unable to satisfy target for token {}: {:?}",
                        self.label.as_str(),
                        e
                    );
                    Ok(None)
                })
            }
        }
    }

    fn satisfy_target(&self) -> Result<Option<PostSyncAction>, Error> {
        let (key_id, key_label) = self
            .db
            .to_schema_keys()
            .iter()
            .map(|k| (k.id, k.label.clone()))
            .next()
            .ok_or(Error::SourceDbEmpty)?;
        let key_uri = format!(
            "pkcs11:model={};manufacturer={};token={};id=%{:02x};object={};type=private?pin-value={}",
            "p11ne-token",
            "Amazon",
            self.label.as_str(),
            key_id,
            key_label.as_str(),
            self.pin.as_str(),
        );

        match self.target {
            None => Ok(None),
            Some(config::Target::Conf {
                ref user,
                ref group,
                ref path,
            }) => {
                let uid = match user.as_ref() {
                    Some(name) => Some(
                        unistd::User::from_name(name)
                            .map_err(Error::NixError)?
                            .ok_or(Error::BadUser(name.to_string()))?
                            .uid,
                    ),
                    None => None,
                };
                let gid = match group.as_ref() {
                    Some(name) => Some(
                        unistd::Group::from_name(name)
                            .map_err(Error::NixError)?
                            .ok_or(Error::BadGroup(name.to_string()))?
                            .gid,
                    ),
                    None => None,
                };

                let cert_path = match self.db.cert_pem() {
                    None => None,
                    Some(cert_pem) => {
                        let cert_name = match self.service {
                            ManagedService::Nginx => "nginx-cert",
                            ManagedService::Httpd => "httpd-cert",
                        };
                        let cert_path = format!(
                            "{}/{}-{}.pem",
                            defs::RUN_DIR,
                            cert_name,
                            util::bytes_to_hex(self.label.as_bytes())
                        );
                        OpenOptions::new()
                            .create(true)
                            .truncate(true)
                            .write(true)
                            .open(cert_path.as_str())
                            .and_then(|mut file| {
                                file.write_all(cert_pem.as_bytes())?;
                                if let Some(chain_pem) = self.db.cert_chain_pem() {
                                    file.write_all(chain_pem.as_bytes())?;
                                }
                                Ok(())
                            })
                            .map_err(Error::TargetIoError)?;
                        Some(cert_path)
                    }
                };

                util::create_dirs_for_file(path).map_err(Error::TargetIoError)?;

                let post_action = match self.service {
                    ManagedService::Nginx => {
                        nginx::NginxService::write_tls_entries(
                            &path, uid, gid, &key_uri, cert_path,
                        )?;
                        Some(PostSyncAction::ReloadNginx)
                    }
                    ManagedService::Httpd => {
                        httpd::HttpdService::write_tls_entries(
                            &path, uid, gid, &key_uri, cert_path,
                        )?;
                        Some(PostSyncAction::ReloadHttpd)
                    }
                };
                Ok(post_action)
            }
        }
    }

    fn to_schema_token(&self) -> Result<schema::Token, Error> {
        Ok(schema::Token {
            label: self.label.clone(),
            pin: self.pin.clone(),
            envelope_key: Self::kms_envelope_key()?,
            keys: self.db.to_schema_keys(),
        })
    }

    fn kms_envelope_key() -> Result<schema::EnvelopeKey, Error> {
        let creds = imds::creds().map_err(Error::ImdsError)?;
        Ok(schema::EnvelopeKey::Kms {
            region: imds::region().map_err(Error::ImdsError)?,
            access_key_id: creds.access_key_id,
            secret_access_key: creds.secret_access_key,
            session_token: creds.token,
        })
    }
}


// Copyright 2020-2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
mod httpd;
mod mngtok;
mod nginx;

use super::defs;
use crate::config;
use crate::gdata;
use crate::imds;
use crate::util;
use crate::{enclave, enclave::P11neEnclave};
use log::{debug, error, info, warn};
use mngtok::{ManagedService, ManagedToken};
use nix::{sys::signal, unistd};
use std::collections::HashSet;
use std::fs::{create_dir_all, OpenOptions};
use std::io::Write;
use std::path::Path;
use std::process::Command;
use std::rc::Rc;
use std::sync::atomic::Ordering;
use std::time::{Duration, Instant};
use vtok_rpc::api::schema;

#[derive(Debug)]
pub enum Error {
    EnclaveBootTimeout,
    EnclaveDied,
    EnclaveError(enclave::Error),
    RemoveTokenError(schema::ApiError),
    ImdsError(imds::Error),
    SendSignalError(nix::Error),
    SystemdExecError(std::io::Error),
    SystemdParsePidError,
    SystemdShowPidError(Option<i32>, String),
    SystemdStartNginxError(Option<i32>),
    SystemdStartHttpdError(Option<i32>),
    Utf8Error(std::string::FromUtf8Error),
    SystemdOverrideError,
    SystemdReloadError,
    TokenError(mngtok::Error),
}

pub struct Agent {
    enclave: Rc<P11neEnclave>,
    tokens: Vec<ManagedToken>,
    options: config::Options,
}

#[derive(Hash, Eq, PartialEq)]
pub enum PostSyncAction {
    ReloadNginx,
    ReloadHttpd,
}

impl Agent {
    pub fn new(mut config: config::Config) -> Result<Self, Error> {
        let enclave = Rc::new(P11neEnclave::new(config.enclave).map_err(Error::EnclaveError)?);
        let service =
            ManagedService::from_str(&config.options.service).map_err(Error::TokenError)?;
        let tokens = config
            .tokens
            .drain(..)
            .filter_map(
                |conf| match ManagedToken::new(conf, enclave.clone(), service) {
                    Ok(tok) => Some(tok),
                    Err(e) => {
                        error!("Error creating managed token: {:?}", e);
                        None
                    }
                },
            )
            .collect();

        debug!("Global options: {:?}", &config.options);

        Ok(Self {
            enclave,
            tokens,
            options: config.options,
        })
    }

    pub fn run(&mut self) -> Result<(), Error> {
        if !self.enclave.wait_boot() {
            return Err(Error::EnclaveBootTimeout);
        }

        let mut next_sync = Instant::now();
        let sync_interval = Duration::from_secs(self.options.sync_interval_secs);

        loop {
            nix::sys::signal::kill(unistd::Pid::from_raw(self.enclave.pid()), None)
                .map_err(|_| Error::EnclaveDied)?;

            if next_sync <= Instant::now() {
                self.sync()?;
                next_sync += sync_interval;
            }

            util::interruptible_sleep(Duration::from_secs(1)).unwrap_or_default();
            if gdata::EXIT_CONDITION.load(Ordering::SeqCst) {
                return Ok(());
            }
        }
    }

    fn sync(&mut self) -> Result<(), Error> {
        let mut broken_list = Vec::new();
        let mut post_actions = HashSet::new();

        imds::invalidate_cache().map_err(Error::ImdsError)?;

        for tok in self.tokens.iter_mut() {
            info!("Syncing token {}", tok.label.as_str());
            match tok.sync() {
                // TODO: tidy up and be more verbose
                Err(mngtok::Error::AcmDbFetchError(_, _))
                | Err(mngtok::Error::AcmDbParseError(_))
                | Err(mngtok::Error::FileDbError(_))
                | Err(mngtok::Error::FileDbParseError(_)) => {
                    error!("Broken token: {}", tok.label.as_str());
                    broken_list.push(tok.label.clone());
                }
                Err(mngtok::Error::AddTokenError(schema::ApiError::KmsDecryptFailed)) => {
                    error!("Attestation failed for new token {}.", tok.label.as_str());
                    broken_list.push(tok.label.clone());
                }
                Err(mngtok::Error::RefreshTokenError(schema::ApiError::KmsDecryptFailed)) => {
                    error!(
                        "Attestation failed for token {}. Releasing",
                        tok.label.as_str()
                    );
                    broken_list.push(tok.label.clone());
                    self.enclave
                        .remove_token(tok.label.clone(), tok.pin.clone())
                        .map_err(Error::EnclaveError)?
                        .map_err(Error::RemoveTokenError)?;
                }
                Err(e) => return Err(Error::TokenError(e)),
                Ok(None) => (),
                Ok(Some(act)) => {
                    post_actions.insert(act);
                }
            }
        }

        // Run post-sync actions
        post_actions
            .drain()
            .map(|act| act.execute(&self.options))
            .count();

        // Remove broken tokens
        self.tokens
            .retain(|tok| broken_list.iter().find(|l| **l == tok.label).is_none());

        Ok(())
    }
}

impl PostSyncAction {
    fn execute(&self, options: &config::Options) {
        info!(
            "Service: {} | Force_Start: {} | Reload: {} | Sync: {}",
            options.service,
            options.force_start,
            options.reload_wait_ms,
            options.sync_interval_secs
        );
        match self {
            Self::ReloadNginx => Self::reload_nginx(options.force_start, options.reload_wait_ms),
            Self::ReloadHttpd => Self::reload_httpd(options.force_start, options.reload_wait_ms),
        }
    }

    fn reload_nginx(force_start: bool, wait_ms: u64) {
        info!("Reloading NGINX configuration.");
        Command::new("systemctl")
            .args(&["show", "--property=MainPID", "nginx.service"])
            .output()
            .map_err(Error::SystemdExecError)
            .and_then(|output| {
                if !output.status.success() {
                    return Err(Error::SystemdShowPidError(
                        output.status.code(),
                        String::from_utf8_lossy(output.stderr.as_slice()).to_string(),
                    ));
                }
                String::from_utf8(output.stdout)
                    .map_err(Error::Utf8Error)
                    .and_then(|line| {
                        line.as_str()
                            .trim()
                            .rsplit("=")
                            .next()
                            .ok_or(Error::SystemdParsePidError)
                            .and_then(|pid_str| {
                                pid_str
                                    .parse::<i32>()
                                    .map_err(|_| Error::SystemdParsePidError)
                            })
                    })
            })
            .and_then(|pid| match (pid, force_start) {
                (0, true) => {
                    info!("NGINX is not running. Starting it now.");
                    Command::new("systemctl")
                        .args(&["start", "nginx.service"])
                        .status()
                        .map_err(Error::SystemdExecError)
                        .and_then(|status| {
                            if !status.success() {
                                Err(Error::SystemdStartNginxError(status.code()))
                            } else {
                                Ok(())
                            }
                        })
                }
                (0, false) => {
                    warn!(
                        "Unable to reload NGINX: it is not running and 'force_start' option is disabled."
                    );
                    Ok(())
                }
                (pid, _) => {
                    debug!("Sending SIGUSR2 to PID={}", pid);
                    signal::kill(unistd::Pid::from_raw(pid), signal::Signal::SIGUSR2)
                        .map_err(Error::SendSignalError)?;
                    debug!("Sending SIGWINCH to PID={}", pid);
                    signal::kill(unistd::Pid::from_raw(pid), signal::Signal::SIGWINCH)
                        .map_err(Error::SendSignalError)?;
                    debug!("Sleeping to allow NGINX to process live update.");
                    std::thread::sleep(Duration::from_millis(wait_ms));
                    debug!("Sending SIGQUIT to PID={}", pid);
                    signal::kill(unistd::Pid::from_raw(pid), signal::Signal::SIGQUIT)
                        .map_err(Error::SendSignalError)?;
                    Ok(())
                }
            })
            .unwrap_or_else(|err| {
                error!("Unable to reload NGINX: {:?}", err);
            });
    }

    fn reload_httpd(force_start: bool, wait_ms: u64) {
        info!("Reloading HTTPD configuration.");
        Command::new("systemctl")
            .args(&["show", "--property=MainPID", "httpd.service"])
            .output()
            .map_err(Error::SystemdExecError)
            .and_then(|output| {
                if !output.status.success() {
                    return Err(Error::SystemdShowPidError(
                        output.status.code(),
                        String::from_utf8_lossy(output.stderr.as_slice()).to_string(),
                    ));
                }
                String::from_utf8(output.stdout)
                    .map_err(Error::Utf8Error)
                    .and_then(|line| {
                        line.as_str()
                            .trim()
                            .rsplit("=")
                            .next()
                            .ok_or(Error::SystemdParsePidError)
                            .and_then(|pid_str| {
                                pid_str
                                    .parse::<i32>()
                                    .map_err(|_| Error::SystemdParsePidError)
                            })
                    })
            })
            .and_then(|pid| match (pid, force_start) {
                (0, true) => {
                    info!("HTTPD is not running. Starting it now.");

                    if !Path::new(defs::HTTPD_OVERRIDE_FILE).exists() {
                        info!("Overriding HTTPD systemd service file.");
                         if let Err(_) = create_dir_all(defs::HTTPD_OVERRIDE_DIR) {
                            return Err(Error::SystemdOverrideError);
                        }
                        if let Err(_) = OpenOptions::new()
                            .create(true)
                            .truncate(true)
                            .write(true)
                            .open(defs::HTTPD_OVERRIDE_FILE)
                            .and_then(|mut file| {
                                file.write_all(defs::HTTPD_OVERRIDE_DATA.as_bytes())?;
                                Ok(())
                            }) {
                                return Err(Error::SystemdOverrideError);
                            }
                        if let Err(_) = Command::new("systemctl")
                            .args(&["daemon-reload"])
                            .status()
                            .map_err(Error::SystemdExecError)
                            .and_then(|status| {
                                if !status.success() {
                                    return Err(Error::SystemdStartHttpdError(status.code()));
                                } else {
                                    Ok(())
                                }
                            }) {
                                return Err(Error::SystemdReloadError);
                            }
                    }
                    Command::new("systemctl")
                        .args(&["start", "httpd.service"])
                        .status()
                        .map_err(Error::SystemdExecError)
                        .and_then(|status| {
                            if !status.success() {
                                Err(Error::SystemdStartHttpdError(status.code()))
                            } else {
                                Ok(())
                            }
                        })
                }
                (0, false) => {
                    warn!(
                        "Unable to reload HTTPD: it is not running and 'force_start' option is disabled."
                    );
                    Ok(())
                }
                (pid, _) => {
                    debug!("Sending SIGWINCH to PID={}", pid);
                    signal::kill(unistd::Pid::from_raw(pid), signal::Signal::SIGWINCH)
                        .map_err(Error::SendSignalError)?;
                    debug!("Sleeping to allow HTTPD to process in-flight requests.");
                    std::thread::sleep(Duration::from_millis(wait_ms));
                    Command::new("systemctl")
                        .args(&["restart", "httpd.service"])
                        .status()
                        .map_err(Error::SystemdExecError)
                        .and_then(|status| {
                            if !status.success() {
                                Err(Error::SystemdStartHttpdError(status.code()))
                            } else {
                                Ok(())
                            }
                        })
                }
            })
            .unwrap_or_else(|err| {
                error!("Unable to reload HTTPD: {:?}", err);
            });
    }
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use super::mngtok::Error;
use nix::unistd;
use std::fs::OpenOptions;
use std::io::Write;
use std::os::unix::fs::OpenOptionsExt;
use std::os::unix::io::AsRawFd;

pub struct NginxService {}

impl NginxService {
    pub fn write_tls_entries(
        path: &str,
        uid: Option<nix::unistd::Uid>,
        gid: Option<nix::unistd::Gid>,
        key_uri: &str,
        cert_path: Option<String>,
    ) -> Result<(), Error> {
        OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .mode(0o440)
            .open(path)
            .map_err(Error::TargetIoError)
            .and_then(|mut file| {
                unistd::fchown(file.as_raw_fd(), uid, gid).map_err(Error::NixError)?;
                nix::sys::stat::fchmod(
                    file.as_raw_fd(),
                    // Safe becase 0o440 is valid.
                    unsafe { nix::sys::stat::Mode::from_bits_unchecked(0o440) },
                )
                .map_err(Error::NixError)?;
                write!(file, "ssl_certificate_key \"engine:pkcs11:{}\";\n", key_uri)
                    .map_err(Error::TargetIoError)?;
                if let Some(cp) = cert_path {
                    write!(file, "ssl_certificate \"{}\";\n", cp).map_err(Error::TargetIoError)?;
                }
                Ok(())
            })?;
        Ok(())
    }
}


use serde::{Deserialize, Serialize};
use serde_json;
use std::fs::OpenOptions;
use std::io::{BufReader, BufWriter, Read, Seek, SeekFrom, Write};
use std::path::Path;

use crate::defs;
use crate::util;
use crate::util::LockedFile;

#[derive(Debug)]
pub enum Error {
    IoError(std::io::Error),
    ReadOnly,
    SerdeError(serde_json::error::Error),
}

#[derive(Deserialize, Serialize)]
pub struct PrivateKey {
    pub pem: String,
    pub encrypted_pem_b64: String,
    pub id: u8,
    pub label: String,
    pub cert_pem: Option<String>,
}

#[derive(Deserialize, Serialize)]
pub struct Token {
    pub label: String,
    pub private_keys: Vec<PrivateKey>,
    pub pin: String,
    /// Token expiry timestamp, in seconds. When CLOCK_MONOTONIC reaches this value, the token
    /// is no longer usable.
    pub expiry_ts: u64,
}

#[derive(Deserialize, Serialize)]
pub struct Device {
    slots: Vec<Option<Token>>,
}

/// Device config helper, used to safely access the global p11ne config, stored in the file system
/// at `crate::defs::DEVICE_CONFIG_PATH`.
pub struct Config {
    device: Device,
    file: Option<LockedFile>,
}

impl Config {
    /// Create a new / default config file.
    ///
    /// Note: this is not concurrency-safe! It should only be called from a non-concurrent context,
    ///       such as at server initialization time.
    pub fn init_new() -> Result<(), Error> {
        let path = Path::new(defs::DEVICE_CONFIG_PATH);

        // It's safe to unwrap here, since we know that defs::DEVICE_CONFIG_PATH is sane.
        let dir_path = path.parent().unwrap();
        std::fs::create_dir_all(dir_path).map_err(Error::IoError)?;

        let file = OpenOptions::new()
            .create(true)
            .write(true)
            .truncate(true)
            .open(path)
            .map_err(Error::IoError)?;

        let mut device = Device {
            slots: Vec::with_capacity(defs::DEVICE_MAX_SLOTS),
        };
        for _ in 0..defs::DEVICE_MAX_SLOTS {
            device.slots.push(None);
        }

        serde_json::to_writer(BufWriter::new(file), &device).map_err(Error::SerdeError)?;

        Ok(())
    }

    /// Load the config in read-only mode.
    /// The config file lock is held only until the the data is read and parsed.
    pub fn load_ro() -> Result<Self, Error> {
        let mut file = LockedFile::open_ro(defs::DEVICE_CONFIG_PATH).map_err(Error::IoError)?;
        let device = Self::load_device(file.as_mut_file())?;
        Ok(Self { device, file: None })
    }

    /// Load the config in read-write mode.
    /// The config lock is held for the entire lifetime of the returned object.
    pub fn load_rw() -> Result<Self, Error> {
        let mut file = LockedFile::open_rw(defs::DEVICE_CONFIG_PATH).map_err(Error::IoError)?;
        let device = Self::load_device(file.as_mut_file())?;
        Ok(Self {
            device,
            file: Some(file),
        })
    }

    /// Write the config to its backing file.
    pub fn save(&mut self) -> Result<(), Error> {
        let file = self.file.as_mut().ok_or(Error::ReadOnly)?;
        file.seek(SeekFrom::Start(0)).map_err(Error::IoError)?;
        file.set_len(0).map_err(Error::IoError)?;
        let writer = BufWriter::new(file.as_mut_file());
        serde_json::to_writer(writer, &self.device).map_err(Error::SerdeError)?;
        file.flush().map_err(Error::IoError)?;
        Ok(())
    }

    pub fn slots(&self) -> &[Option<Token>] {
        &self.device.slots
    }

    pub fn slots_mut(&mut self) -> &mut [Option<Token>] {
        &mut self.device.slots
    }

    fn load_device<R: Read>(src: R) -> Result<Device, Error> {
        let mut device: Device =
            serde_json::from_reader(BufReader::new(src)).map_err(Error::SerdeError)?;
        for slot in device.slots.iter_mut() {
            let expired = slot
                .as_ref()
                .and_then(|tok| tok.expiry_ts.checked_sub(util::time::monotonic_secs()))
                .filter(|t| *t > 0)
                .is_none();
            if expired {
                slot.take();
            }
        }
        Ok(device)
    }
}


// Copyright 2020-2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

extern crate libc;

pub mod config;
pub mod util;

pub mod defs {
    /// Manufacturer of the various PKCS#11 objects (library, tokens, slots).
    pub const MANUFACTURER: &str = "Amazon";

    /// The filesystem path used to store the p11ne DB/config file. This is read by the PKCS#11
    /// provider (vtok_p11), and written to by the RPC server (vtok_srv).
    pub const DEVICE_CONFIG_PATH: &str = "./config.json";
    /// PKCS#11 provider description, as reported by vtok_p11.
    pub const DEVICE_DESCRIPTION: &str = "p11ne";
    /// Maximum number of slots/tokens exposed by our PKCS#11 provider.
    pub const DEVICE_MAX_SLOTS: usize = 128;

    /// PKCS#11 provider slot description.
    pub const SLOT_DESCRIPTION: &str = "p11ne-slot";

    /// Token expiration time, in seconds. A refresh operation (i.e. reattestation) MUST be
    /// performed before this timer expires. Otherwise, the token will be automatically removed
    /// from its slot.
    pub const TOKEN_EXPIRY_SECS: u64 = 48 * 3600;
    /// PKCS#11 token model, as reported by vtok_p11.
    pub const TOKEN_MODEL: &str = "p11ne-token";
    /// Minimum length (in bytes) of a token PIN.
    pub const TOKEN_MIN_PIN_LEN: usize = 4;
    /// Maximum length (in bytes) of a token PIN.
    pub const TOKEN_MAX_PIN_LEN: usize = 64;
    /// Minimum length (in bytes) of a token label.
    pub const TOKEN_MIN_LABEL_LEN: usize = 1;
    /// Maximum length (in bytes) of a token label.
    pub const TOKEN_MAX_LABEL_LEN: usize = 32;
    /// Maximum number of concurrent (read-only) sessions supported by a single token exposed by
    /// vtok_p11.
    pub const TOKEN_MAX_SESSIONS: usize = 1024;

    /// I/O timeout (in milliseconds) set for the RPC client-server communication.
    pub const RPC_STREAM_TIMEOUT_MS: u64 = 1000;
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use libc;
use std::fs::{File, OpenOptions};
use std::io::{Error as IoError, ErrorKind, Result as IoResult};
use std::ops::{Deref, DerefMut};
use std::os::unix::io::AsRawFd;
use std::path::Path;

/// Wrap an `std::fs::File` object, ensuring that its inner FD is locked
/// via `libc::flock`.
///
/// `Deref` and `DerefMut` are employed in order to easily expose the original
/// `File` interface (i.e. `LockedFile` can be derefed to `File`).
pub struct LockedFile(File);

impl LockedFile {
    /// Create a `LockedFile` from the file at `path`, opened in read-only mode.
    /// Note: the lock itself will be a read-only lock.
    pub fn open_ro<P: AsRef<Path>>(path: P) -> IoResult<Self> {
        OpenOptions::new()
            .read(true)
            .open(path)
            .and_then(|f| Self::from_file(f, libc::LOCK_SH))
    }

    /// Create a `LockedFile` from the file at `path`, opened in read-write mode.
    /// Note: the lock itself will be a read-write (i.e. exclusive) lock.
    pub fn open_rw<P: AsRef<Path>>(path: P) -> IoResult<Self> {
        OpenOptions::new()
            .read(true)
            .write(true)
            .open(path)
            .and_then(|f| Self::from_file(f, libc::LOCK_EX))
    }

    /// Get a mutable reference to the inner `File` object.
    pub fn as_mut_file(&mut self) -> &mut File {
        &mut self.0
    }

    fn from_file(file: File, flock_op: libc::c_int) -> IoResult<Self> {
        loop {
            let rc = unsafe { libc::flock(file.as_raw_fd(), flock_op) };
            if rc == 0 {
                break;
            }
            let err = IoError::last_os_error();

            // If our wait was interrupted, try to acquire the lock again.
            if err.kind() != ErrorKind::Interrupted {
                return Err(err);
            }
        }
        Ok(Self(file))
    }
}

impl Deref for LockedFile {
    type Target = File;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for LockedFile {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl Drop for LockedFile {
    fn drop(&mut self) {
        unsafe {
            libc::flock(self.0.as_raw_fd(), libc::LOCK_UN);
        }
    }
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

pub mod flock;
pub mod time;

pub use flock::LockedFile;


use libc;

/// Get a monotonic (i.e. always increasing) timestamp, in seconds.
pub fn monotonic_secs() -> u64 {
    let mut ts = libc::timespec {
        tv_sec: 0,
        tv_nsec: 0,
    };

    let rc = unsafe { libc::clock_gettime(libc::CLOCK_MONOTONIC, &mut ts) };
    if rc != 0 {
        panic!("clock_gettime() failed: {}", rc);
    }
    if ts.tv_sec < 0 {
        panic!("Got negative value for CLOCK_MONOTONIC: {}", ts.tv_sec);
    }

    ts.tv_sec as u64
}


// Copyright 2020-2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use std::process::Command;

/// The p11ne enclave init
/// Spawns the p11ne enclave main applications
/// - the provisioning/rpc server
/// - the p11-kit server
fn main() {
    Command::new("p11-kit")
        .args(&[
            "server",
            "-n",
            "vsock:port=9999",
            "--provider",
            "/usr/lib/libvtok_p11.so",
            "-f",
            "-v",
            "pkcs11:",
        ])
        .spawn()
        .expect("p11-kit server failed to start.");
    Command::new("p11ne-server")
        .args(&["vsock", "10000"])
        .spawn()
        .expect("provisioning server failed to start.")
        .wait() // Block here. If the provisioning server dies, terminate the enclave.
        .expect("provisioning server has exited.");
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

fn main() {
    println!("cargo:rustc-link-lib=dylib=crypto");
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

extern crate lazy_static;
#[macro_use]
extern crate log;
extern crate serde;
extern crate serde_json;
extern crate vtok_common;

#[macro_use]
mod util;
mod api;
mod backend;
mod crypto;
mod pkcs11;

use crate::crypto::Error as CryptoError;
use backend::token::Error as TokenError;

/// Device, slot and token capabilities and information
mod defs {
    use super::backend::mech::{MechDigest, Mechanism};
    use super::pkcs11;
    use vtok_common::defs as comm_defs;

    pub const CRYPTOKI_VERSION: pkcs11::CK_VERSION = ck_version!(
        pkcs11::CRYPTOKI_VERSION_MAJOR as u8,
        pkcs11::CRYPTOKI_VERSION_MINOR as u8
    );

    pub const DEVICE_DESCRIPTION: &str = comm_defs::DEVICE_DESCRIPTION;
    pub const DEVICE_VERSION: pkcs11::CK_VERSION = ck_version!(0, 1);
    pub const DEVICE_MANUFACTURER: &str = comm_defs::MANUFACTURER;

    pub const SLOT_DESCRIPTION: &str = comm_defs::SLOT_DESCRIPTION;
    pub const SLOT_HARDWARE_VERSION: pkcs11::CK_VERSION = ck_version!(0, 1);
    pub const SLOT_FIRMWARE_VERSION: pkcs11::CK_VERSION = ck_version!(0, 1);
    pub const SLOT_MANUFACTURER: &str = comm_defs::MANUFACTURER;
    pub const MAX_SLOTS: usize = comm_defs::DEVICE_MAX_SLOTS;

    pub const TOKEN_MANUFACTURER: &str = comm_defs::MANUFACTURER;
    pub const TOKEN_MODEL: &str = comm_defs::TOKEN_MODEL;
    pub const TOKEN_MAX_SESSIONS: pkcs11::CK_ULONG =
        comm_defs::TOKEN_MAX_SESSIONS as pkcs11::CK_ULONG;
    pub const TOKEN_MAX_RW_SESSIONS: pkcs11::CK_ULONG = 0;
    pub const TOKEN_MAX_PIN_LEN: pkcs11::CK_ULONG =
        comm_defs::TOKEN_MIN_PIN_LEN as pkcs11::CK_ULONG;
    pub const TOKEN_MIN_PIN_LEN: pkcs11::CK_ULONG =
        comm_defs::TOKEN_MAX_PIN_LEN as pkcs11::CK_ULONG;
    pub const TOKEN_HARDWARE_VERSION: pkcs11::CK_VERSION = ck_version!(0, 1);
    pub const TOKEN_FIRMWARE_VERSION: pkcs11::CK_VERSION = ck_version!(0, 1);
    pub const TOKEN_UTC_TIME: &str = "";

    pub const TOKEN_MECH_LIST: [Mechanism; 24] = [
        Mechanism::Digest(MechDigest::Sha1),
        Mechanism::Digest(MechDigest::Sha224),
        Mechanism::Digest(MechDigest::Sha256),
        Mechanism::Digest(MechDigest::Sha384),
        Mechanism::Digest(MechDigest::Sha512),
        Mechanism::RsaPkcs(None),
        Mechanism::RsaPkcs(Some(MechDigest::Sha1)),
        Mechanism::RsaPkcs(Some(MechDigest::Sha224)),
        Mechanism::RsaPkcs(Some(MechDigest::Sha256)),
        Mechanism::RsaPkcs(Some(MechDigest::Sha384)),
        Mechanism::RsaPkcs(Some(MechDigest::Sha512)),
        Mechanism::RsaPkcsPss(None, None),
        Mechanism::RsaPkcsPss(Some(MechDigest::Sha1), None),
        Mechanism::RsaPkcsPss(Some(MechDigest::Sha224), None),
        Mechanism::RsaPkcsPss(Some(MechDigest::Sha256), None),
        Mechanism::RsaPkcsPss(Some(MechDigest::Sha384), None),
        Mechanism::RsaPkcsPss(Some(MechDigest::Sha512), None),
        Mechanism::RsaX509,
        Mechanism::Ecdsa(None),
        Mechanism::Ecdsa(Some(MechDigest::Sha1)),
        Mechanism::Ecdsa(Some(MechDigest::Sha224)),
        Mechanism::Ecdsa(Some(MechDigest::Sha256)),
        Mechanism::Ecdsa(Some(MechDigest::Sha384)),
        Mechanism::Ecdsa(Some(MechDigest::Sha512)),
    ];
}

/// Helper for comparing Criptoki versions
impl PartialEq for pkcs11::CK_VERSION {
    fn eq(&self, other: &Self) -> bool {
        self.major == other.major && self.minor == other.minor
    }
}

mod data {
    use lazy_static::lazy_static;
    use std::sync::{Arc, Mutex};

    use super::api;
    use super::api::nyi;
    use super::backend::device::Device;
    use super::defs;
    use super::pkcs11::*;

    /// Supported Criptoki PKCS#11 functions. As per specification,
    /// all functions are exported but some return CKR_FUNCTION_NOT_SUPPORTED.
    pub static FN_LIST: CK_FUNCTION_LIST = CK_FUNCTION_LIST {
        version: defs::DEVICE_VERSION,
        C_Initialize: Some(api::C_Initialize),
        C_Finalize: Some(api::C_Finalize),
        C_GetInfo: Some(api::C_GetInfo),
        C_GetFunctionList: Some(api::C_GetFunctionList),
        C_GetSlotList: Some(api::token::C_GetSlotList),
        C_GetSlotInfo: Some(api::token::C_GetSlotInfo),
        C_GetTokenInfo: Some(api::token::C_GetTokenInfo),
        C_GetMechanismList: Some(api::token::C_GetMechanismList),
        C_GetMechanismInfo: Some(api::token::C_GetMechanismInfo),
        C_InitToken: Some(api::token::C_InitToken),
        C_InitPIN: Some(nyi::C_InitPIN),
        C_SetPIN: Some(nyi::C_SetPIN),
        C_OpenSession: Some(api::session::C_OpenSession),
        C_CloseSession: Some(api::session::C_CloseSession),
        C_CloseAllSessions: Some(api::session::C_CloseAllSessions),
        C_GetSessionInfo: Some(api::session::C_GetSessionInfo),
        C_GetOperationState: Some(nyi::C_GetOperationState),
        C_SetOperationState: Some(nyi::C_SetOperationState),
        C_Login: Some(api::token::C_Login),
        C_Logout: Some(api::token::C_Logout),
        C_CreateObject: Some(nyi::C_CreateObject),
        C_CopyObject: Some(nyi::C_CopyObject),
        C_DestroyObject: Some(nyi::C_DestroyObject),
        C_GetObjectSize: Some(api::object::C_GetObjectSize),
        C_GetAttributeValue: Some(api::object::C_GetAttributeValue),
        C_SetAttributeValue: Some(nyi::C_SetAttributeValue),
        C_FindObjectsInit: Some(api::object::C_FindObjectsInit),
        C_FindObjects: Some(api::object::C_FindObjects),
        C_FindObjectsFinal: Some(api::object::C_FindObjectsFinal),
        C_EncryptInit: Some(api::encrypt::C_EncryptInit),
        C_Encrypt: Some(api::encrypt::C_Encrypt),
        C_EncryptUpdate: Some(nyi::C_EncryptUpdate),
        C_EncryptFinal: Some(nyi::C_EncryptFinal),
        C_DecryptInit: Some(api::decrypt::C_DecryptInit),
        C_Decrypt: Some(api::decrypt::C_Decrypt),
        C_DecryptUpdate: Some(nyi::C_DecryptUpdate),
        C_DecryptFinal: Some(nyi::C_DecryptFinal),
        C_DigestInit: Some(api::digest::C_DigestInit),
        C_Digest: Some(api::digest::C_Digest),
        C_DigestUpdate: Some(api::digest::C_DigestUpdate),
        C_DigestKey: Some(nyi::C_DigestKey),
        C_DigestFinal: Some(api::digest::C_DigestFinal),
        C_SignInit: Some(api::sign::C_SignInit),
        C_Sign: Some(api::sign::C_Sign),
        C_SignUpdate: Some(api::sign::C_SignUpdate),
        C_SignFinal: Some(api::sign::C_SignFinal),
        C_SignRecoverInit: Some(nyi::C_SignRecoverInit),
        C_SignRecover: Some(nyi::C_SignRecover),
        C_VerifyInit: Some(api::verify::C_VerifyInit),
        C_Verify: Some(api::verify::C_Verify),
        C_VerifyUpdate: Some(api::verify::C_VerifyUpdate),
        C_VerifyFinal: Some(api::verify::C_VerifyFinal),
        C_VerifyRecoverInit: Some(nyi::C_VerifyRecoverInit),
        C_VerifyRecover: Some(nyi::C_VerifyRecover),
        C_DigestEncryptUpdate: Some(nyi::C_DigestEncryptUpdate),
        C_DecryptDigestUpdate: Some(nyi::C_DecryptDigestUpdate),
        C_SignEncryptUpdate: Some(nyi::C_SignEncryptUpdate),
        C_DecryptVerifyUpdate: Some(nyi::C_DecryptVerifyUpdate),
        C_GenerateKey: Some(nyi::C_GenerateKey),
        C_GenerateKeyPair: Some(nyi::C_GenerateKeyPair),
        C_WrapKey: Some(nyi::C_WrapKey),
        C_UnwrapKey: Some(nyi::C_UnwrapKey),
        C_DeriveKey: Some(nyi::C_DeriveKey),
        C_SeedRandom: Some(nyi::C_SeedRandom),
        C_GenerateRandom: Some(nyi::C_GenerateRandom),
        C_GetFunctionStatus: Some(api::session::C_GetFunctionStatus),
        C_CancelFunction: Some(api::session::C_CancelFunction),
        C_WaitForSlotEvent: Some(nyi::C_WaitForSlotEvent),
    };

    lazy_static! {
        pub static ref DEVICE: Arc<Mutex<Option<Device>>> = Arc::new(Mutex::new(None));
    }
}

#[derive(Clone, Copy, Debug)]
pub enum Error {
    CkError(pkcs11::CK_RV),
    CryptoError(crypto::Error),
    KeyHandleInvalid,
    KeyTypeInconsistent,
    MechanismInvalid,
    OperationNotInitialized,
    SessionHandleInvalid,
    SessionLockPoisoned,
    SlotIdInvalid,
    UserNotLoggedIn,
    TokenError(TokenError),
    TokenNotPresent,
    TokenUninit,
    // TODO: implement proper error reporting
    GeneralError,
}

impl Error {
    pub fn ck_rv(&self) -> pkcs11::CK_RV {
        match self {
            Self::CkError(code) => *code,
            Self::CryptoError(CryptoError::DataMissing) => pkcs11::CKR_DATA_LEN_RANGE,
            Self::CryptoError(CryptoError::OperationActive) => pkcs11::CKR_OPERATION_ACTIVE,
            Self::CryptoError(CryptoError::DigestVerifyFinal) => pkcs11::CKR_SIGNATURE_INVALID,
            Self::CryptoError(CryptoError::DigestVerify) => pkcs11::CKR_SIGNATURE_INVALID,
            Self::CryptoError(CryptoError::DirectVerify) => pkcs11::CKR_SIGNATURE_INVALID,
            Self::KeyHandleInvalid => pkcs11::CKR_KEY_HANDLE_INVALID,
            Self::KeyTypeInconsistent => pkcs11::CKR_KEY_TYPE_INCONSISTENT,
            Self::MechanismInvalid => pkcs11::CKR_MECHANISM_INVALID,
            Self::OperationNotInitialized => pkcs11::CKR_OPERATION_NOT_INITIALIZED,
            Self::SessionHandleInvalid => pkcs11::CKR_SESSION_HANDLE_INVALID,
            Self::SlotIdInvalid => pkcs11::CKR_SLOT_ID_INVALID,
            Self::UserNotLoggedIn => pkcs11::CKR_USER_NOT_LOGGED_IN,
            Self::TokenNotPresent => pkcs11::CKR_TOKEN_NOT_PRESENT,
            Self::TokenError(TokenError::SessionCount) => pkcs11::CKR_SESSION_COUNT,
            Self::TokenError(TokenError::SessionHandleInvalid) => {
                pkcs11::CKR_SESSION_HANDLE_INVALID
            }
            Self::TokenError(TokenError::UserAlreadyLoggedIn) => pkcs11::CKR_USER_ALREADY_LOGGED_IN,
            Self::TokenError(TokenError::UserNotLoggedIn) => pkcs11::CKR_USER_NOT_LOGGED_IN,
            Self::TokenError(TokenError::MechNotFound) => pkcs11::CKR_MECHANISM_INVALID,
            Self::TokenError(TokenError::PinIncorrect) => pkcs11::CKR_PIN_INCORRECT,
            _ => pkcs11::CKR_GENERAL_ERROR,
        }
    }
}

impl From<Error> for pkcs11::CK_RV {
    fn from(src: Error) -> pkcs11::CK_RV {
        src.ck_rv()
    }
}

pub type Result<T> = std::result::Result<T, Error>;