pub mod command_parser;
pub mod protocol_helpers;
pub mod utils;

use command_parser::{CommandOutput, FileArgs, ListenArgs, RunArgs};
use protocol_helpers::{recv_loop, recv_u64, send_loop, send_u64};

use nix::sys::socket::listen as listen_vsock;
use nix::sys::socket::{accept, bind, connect, shutdown, socket};
use nix::sys::socket::{AddressFamily, Shutdown, SockAddr, SockFlag, SockType};
use nix::unistd::close;
use num_derive::FromPrimitive;
use num_traits::FromPrimitive;
use std::cmp::min;
use std::convert::TryInto;
use std::fs::File;
use std::io::{Read, Write};
use std::os::unix::io::{AsRawFd, RawFd};
use std::process::Command;

pub const VMADDR_CID_ANY: u32 = 0xFFFFFFFF;
pub const BUF_MAX_LEN: usize = 8192;
pub const BACKLOG: usize = 128;
const MAX_CONNECTION_ATTEMPTS: usize = 5;

#[derive(Debug, Clone, FromPrimitive)]
enum CmdId {
    RunCmd = 0,
    RecvFile,
    SendFile,
    RunCmdNoWait,
}

struct VsockSocket {
    socket_fd: RawFd,
}

impl VsockSocket {
    fn new(socket_fd: RawFd) -> Self {
        VsockSocket { socket_fd }
    }
}

impl Drop for VsockSocket {
    fn drop(&mut self) {
        shutdown(self.socket_fd, Shutdown::Both)
            .unwrap_or_else(|e| eprintln!("Failed to shut socket down: {:?}", e));
        close(self.socket_fd).unwrap_or_else(|e| eprintln!("Failed to close socket: {:?}", e));
    }
}

impl AsRawFd for VsockSocket {
    fn as_raw_fd(&self) -> RawFd {
        self.socket_fd
    }
}

fn vsock_connect(cid: u32, port: u32) -> Result<VsockSocket, String> {
    let sockaddr = SockAddr::new_vsock(cid, port);
    let mut err_msg = String::new();

    for i in 0..MAX_CONNECTION_ATTEMPTS {
        let vsocket = VsockSocket::new(
            socket(
                AddressFamily::Vsock,
                SockType::Stream,
                SockFlag::empty(),
                None,
            )
            .map_err(|err| format!("Failed to create the socket: {:?}", err))?,
        );
        match connect(vsocket.as_raw_fd(), &sockaddr) {
            Ok(_) => return Ok(vsocket),
            Err(e) => err_msg = format!("Failed to connect: {}", e),
        }

        std::thread::sleep(std::time::Duration::from_secs(1 << i));
    }

    Err(err_msg)
}

fn run_server(fd: RawFd, no_wait: bool) -> Result<(), String> {
    // recv command
    let len = recv_u64(fd)?;
    let mut buf = [0u8; BUF_MAX_LEN];
    recv_loop(fd, &mut buf, len)?;

    let len_usize = len.try_into().map_err(|err| format!("{:?}", err))?;
    let command = std::str::from_utf8(&buf[0..len_usize]).map_err(|err| format!("{:?}", err))?;

    // execute command
    let command_output = if no_wait {
        #[rustfmt::skip]
        let output = Command::new("sh")
            .arg("-c")
            .arg(command)
            .spawn();
        if output.is_err() {
            CommandOutput::new(
                String::new(),
                format!("Could not execute the command {}", command),
                1,
            )
        } else {
            CommandOutput::new(String::new(), String::new(), 0)
        }
    } else {
        let output = Command::new("sh")
            .arg("-c")
            .arg(command)
            .output()
            .map_err(|err| format!("Could not execute the command {}: {:?}", command, err))?;
        CommandOutput::new_from(output)?
    };

    // send output
    let json_output = serde_json::to_string(&command_output)
        .map_err(|err| format!("Could not serialize the output: {:?}", err))?;
    let buf = json_output.as_bytes();
    let len: u64 = buf.len().try_into().map_err(|err| format!("{:?}", err))?;
    send_u64(fd, len)?;
    send_loop(fd, &buf, len)?;
    Ok(())
}

fn recv_file_server(fd: RawFd) -> Result<(), String> {
    // recv file path
    let len = recv_u64(fd)?;
    let mut buf = [0u8; BUF_MAX_LEN];
    recv_loop(fd, &mut buf, len)?;
    let len_usize = len.try_into().map_err(|err| format!("{:?}", err))?;
    let path = std::str::from_utf8(&buf[0..len_usize]).map_err(|err| format!("{:?}", err))?;

    let mut file = File::open(path).map_err(|err| format!("Could not open file {:?}", err))?;

    let filesize = file
        .metadata()
        .map_err(|err| format!("Could not get file metadata {:?}", err))?
        .len();

    send_u64(fd, filesize)?;
    println!("Sending file {} - size {}", path, filesize);

    let mut progress: u64 = 0;
    let mut tmpsize: u64;

    while progress < filesize {
        tmpsize = buf.len().try_into().map_err(|err| format!("{:?}", err))?;
        tmpsize = min(tmpsize, filesize - progress);

        file.read_exact(&mut buf[..tmpsize.try_into().map_err(|err| format!("{:?}", err))?])
            .map_err(|err| format!("Could not read {:?}", err))?;
        send_loop(fd, &buf, tmpsize)?;
        progress += tmpsize
    }

    Ok(())
}

fn send_file_server(fd: RawFd) -> Result<(), String> {
    // recv file path
    let len = recv_u64(fd)?;
    let mut buf = [0u8; BUF_MAX_LEN];
    recv_loop(fd, &mut buf, len)?;
    let len_usize = len.try_into().map_err(|err| format!("{:?}", err))?;
    let path = std::str::from_utf8(&buf[0..len_usize]).map_err(|err| format!("{:?}", err))?;

    let mut file = File::create(path).map_err(|err| format!("Could not open file {:?}", err))?;

    // Receive filesize
    let filesize = recv_u64(fd)?;
    println!("Receiving file {} - size {}", path, filesize);

    let mut progress: u64 = 0;
    let mut tmpsize: u64;

    while progress < filesize {
        tmpsize = buf.len().try_into().map_err(|err| format!("{:?}", err))?;
        tmpsize = min(tmpsize, filesize - progress);

        recv_loop(fd, &mut buf, tmpsize)?;
        file.write_all(&buf[..tmpsize.try_into().map_err(|err| format!("{:?}", err))?])
            .map_err(|err| format!("Could not write {:?}", err))?;
        progress += tmpsize
    }

    Ok(())
}

pub fn listen(args: ListenArgs) -> Result<(), String> {
    let socket_fd = socket(
        AddressFamily::Vsock,
        SockType::Stream,
        SockFlag::empty(),
        None,
    )
    .map_err(|err| format!("Create socket failed: {:?}", err))?;

    let sockaddr = SockAddr::new_vsock(VMADDR_CID_ANY, args.port);

    bind(socket_fd, &sockaddr).map_err(|err| format!("Bind failed: {:?}", err))?;

    listen_vsock(socket_fd, BACKLOG).map_err(|err| format!("Listen failed: {:?}", err))?;

    loop {
        let fd = accept(socket_fd).map_err(|err| format!("Accept failed: {:?}", err))?;

        //cmd id
        let cmdid = match recv_u64(fd) {
            Ok(id_u64) => match CmdId::from_u64(id_u64) {
                Some(c) => c,
                _ => {
                    eprintln!("Error no such command");
                    continue;
                }
            },
            Err(e) => {
                eprintln!("Error {}", e);
                continue;
            }
        };

        match cmdid {
            CmdId::RunCmd => {
                if let Err(e) = run_server(fd, false) {
                    eprintln!("Error {}", e);
                }
            }
            CmdId::RecvFile => {
                if let Err(e) = recv_file_server(fd) {
                    eprintln!("Error {}", e);
                }
            }
            CmdId::SendFile => {
                if let Err(e) = send_file_server(fd) {
                    eprintln!("Error {}", e);
                }
            }
            CmdId::RunCmdNoWait => {
                if let Err(e) = run_server(fd, true) {
                    eprintln!("Error {}", e);
                }
            }
        }
    }
}

pub fn run(args: RunArgs) -> Result<i32, String> {
    let vsocket = vsock_connect(args.cid, args.port)?;
    let socket_fd = vsocket.as_raw_fd();

    // Send command id
    if args.no_wait {
        send_u64(socket_fd, CmdId::RunCmdNoWait as u64)?;
    } else {
        send_u64(socket_fd, CmdId::RunCmd as u64)?;
    }

    // send command
    let buf = args.command.as_bytes();
    let len: u64 = buf.len().try_into().map_err(|err| format!("{:?}", err))?;
    send_u64(socket_fd, len)?;
    send_loop(socket_fd, &buf, len)?;

    // recv output
    let mut buf = [0u8; BUF_MAX_LEN];
    let len = recv_u64(socket_fd)?;
    let mut json_output = String::new();
    let mut to_recv = len;
    while to_recv > 0 {
        let recv_len = min(BUF_MAX_LEN as u64, to_recv);
        recv_loop(socket_fd, &mut buf, recv_len)?;
        to_recv -= recv_len;
        let to_recv_usize: usize = recv_len.try_into().map_err(|err| format!("{:?}", err))?;
        json_output.push_str(
            std::str::from_utf8(&buf[0..to_recv_usize]).map_err(|err| format!("{:?}", err))?,
        );
    }

    let output: CommandOutput = serde_json::from_str(json_output.as_str())
        .map_err(|err| format!("Could not deserialize the output: {:?}", err))?;
    print!("{}", output.stdout);
    eprint!("{}", output.stderr);

    let rc = match output.rc {
        Some(code) => code,
        _ => 0,
    };

    Ok(rc)
}

pub fn recv_file(args: FileArgs) -> Result<(), String> {
    let mut file = File::create(&args.localfile)
        .map_err(|err| format!("Could not open localfile {:?}", err))?;
    let vsocket = vsock_connect(args.cid, args.port)?;
    let socket_fd = vsocket.as_raw_fd();

    // Send command id
    send_u64(socket_fd, CmdId::RecvFile as u64)?;

    // send remotefile path
    let buf = args.remotefile.as_bytes();
    let len: u64 = buf.len().try_into().map_err(|err| format!("{:?}", err))?;
    send_u64(socket_fd, len)?;
    send_loop(socket_fd, &buf, len)?;

    // Receive filesize
    let filesize = recv_u64(socket_fd)?;
    println!(
        "Receiving file {}(saving to {}) - size {}",
        &args.remotefile,
        &args.localfile[..],
        filesize
    );

    let mut progress: u64 = 0;
    let mut tmpsize: u64;
    let mut buf = [0u8; BUF_MAX_LEN];

    while progress < filesize {
        tmpsize = buf.len().try_into().map_err(|err| format!("{:?}", err))?;
        tmpsize = min(tmpsize, filesize - progress);

        recv_loop(socket_fd, &mut buf, tmpsize)?;
        file.write_all(&buf[..tmpsize.try_into().map_err(|err| format!("{:?}", err))?])
            .map_err(|err| format!("Could not write {:?}", err))?;
        progress += tmpsize
    }
    Ok(())
}

pub fn send_file(args: FileArgs) -> Result<(), String> {
    let mut file =
        File::open(&args.localfile).map_err(|err| format!("Could not open localfile {:?}", err))?;
    let vsocket = vsock_connect(args.cid, args.port)?;
    let socket_fd = vsocket.as_raw_fd();

    // Send command id
    send_u64(socket_fd, CmdId::SendFile as u64)?;

    // send remotefile path
    let buf = args.remotefile.as_bytes();
    let len: u64 = buf.len().try_into().map_err(|err| format!("{:?}", err))?;
    send_u64(socket_fd, len)?;
    send_loop(socket_fd, &buf, len)?;

    let filesize = file
        .metadata()
        .map_err(|err| format!("Could not get file metadate {:?}", err))?
        .len();

    send_u64(socket_fd, filesize)?;
    println!(
        "Sending file {}(sending to {}) - size {}",
        &args.localfile,
        &args.remotefile[..],
        filesize
    );

    let mut buf = [0u8; BUF_MAX_LEN];
    let mut progress: u64 = 0;
    let mut tmpsize: u64;

    while progress < filesize {
        tmpsize = buf.len().try_into().map_err(|err| format!("{:?}", err))?;
        tmpsize = min(tmpsize, filesize - progress);

        file.read_exact(&mut buf[..tmpsize.try_into().map_err(|err| format!("{:?}", err))?])
            .map_err(|err| format!("Could not read {:?}", err))?;
        send_loop(socket_fd, &buf, tmpsize)?;
        progress += tmpsize
    }

    Ok(())
}


use clap::{App, AppSettings, Arg, SubCommand};

use command_executer::command_parser::{FileArgs, ListenArgs, RunArgs};
use command_executer::create_app;
use command_executer::{listen, recv_file, run, send_file};

fn main() {
    let app = create_app!();
    let args = app.get_matches();

    match args.subcommand() {
        Some(("listen", args)) => {
            let listen_args = ListenArgs::new_with(args).unwrap();
            listen(listen_args).unwrap();
        }
        Some(("run", args)) => {
            let run_args = RunArgs::new_with(args).unwrap();
            let rc = run(run_args).unwrap();
            std::process::exit(rc);
        }
        Some(("recv-file", args)) => {
            let subcmd_args = FileArgs::new_with(args).unwrap();
            recv_file(subcmd_args).unwrap();
        }
        Some(("send-file", args)) => {
            let subcmd_args = FileArgs::new_with(args).unwrap();
            send_file(subcmd_args).unwrap();
        }
        Some(_) | None => {}
    }
}


use byteorder::{ByteOrder, LittleEndian};
use nix::sys::socket::MsgFlags;
use nix::sys::socket::{recv, send};
use std::convert::TryInto;
use std::os::unix::io::RawFd;

pub fn send_u64(fd: RawFd, val: u64) -> Result<(), String> {
    let mut buf = [0u8; 9];
    LittleEndian::write_u64(&mut buf, val);
    send_loop(fd, &mut buf, 9)?;
    Ok(())
}

pub fn recv_u64(fd: RawFd) -> Result<u64, String> {
    let mut buf = [0u8; 9];
    recv_loop(fd, &mut buf, 9)?;
    let val = LittleEndian::read_u64(&buf);
    Ok(val)
}

pub fn send_i32(fd: RawFd, val: i32) -> Result<(), String> {
    let mut buf = [0u8; 4];
    LittleEndian::write_i32(&mut buf, val);
    send_loop(fd, &mut buf, 4)?;
    Ok(())
}

pub fn recv_i32(fd: RawFd) -> Result<i32, String> {
    let mut buf = [0u8; 4];
    recv_loop(fd, &mut buf, 4)?;
    let val = LittleEndian::read_i32(&buf);
    Ok(val)
}

pub fn send_loop(fd: RawFd, buf: &[u8], len: u64) -> Result<(), String> {
    let len: usize = len.try_into().map_err(|err| format!("{:?}", err))?;
    let mut send_bytes = 0;

    while send_bytes < len {
        let size = match send(fd, &buf[send_bytes..len], MsgFlags::empty()) {
            Ok(size) => size,
            Err(nix::errno::Errno::EINTR) => 0,
            Err(err) => return Err(format!("{:?}", err)),
        };
        send_bytes += size;
    }

    Ok(())
}

pub fn recv_loop(fd: RawFd, buf: &mut [u8], len: u64) -> Result<(), String> {
    let len: usize = len.try_into().map_err(|err| format!("{:?}", err))?;
    let mut recv_bytes = 0;

    while recv_bytes < len {
        let size = match recv(fd, &mut buf[recv_bytes..len], MsgFlags::empty()) {
            Ok(size) => size,
            Err(nix::errno::Errno::EINTR) => 0,
            Err(err) => return Err(format!("{:?}", err)),
        };
        recv_bytes += size;
    }

    Ok(())
}


use log::error;

pub trait ExitGracefully<T, E> {
    fn ok_or_exit(self, message: &str) -> T;
}

impl<T, E: std::fmt::Debug> ExitGracefully<T, E> for Result<T, E> {
    fn ok_or_exit(self, message: &str) -> T {
        match self {
            Ok(val) => val,
            Err(err) => {
                error!("{:?}: {}", err, message);
                std::process::exit(1);
            }
        }
    }
}

#[macro_export]
macro_rules! create_app {
    () => {
        App::new("Vsock Tool")
            .about("Tool that runs commands inside the enclave")
            .setting(AppSettings::ArgRequiredElseHelp)
            .version(env!("CARGO_PKG_VERSION"))
            .subcommand(
                SubCommand::with_name("listen")
                    .about("Listen on a given port")
                    .arg(
                        Arg::with_name("port")
                            .long("port")
                            .help("port")
                            .takes_value(true)
                            .required(true),
                    ),
            )
            .subcommand(
                SubCommand::with_name("run")
                    .about("Run a command inside the enclave")
                    .arg(
                        Arg::with_name("port")
                            .long("port")
                            .help("port")
                            .takes_value(true)
                            .required(true),
                    )
                    .arg(
                        Arg::with_name("cid")
                            .long("cid")
                            .help("cid")
                            .takes_value(true)
                            .required(true),
                    )
                    .arg(
                        Arg::with_name("command")
                            .long("command")
                            .help("command")
                            .takes_value(true)
                            .required(true),
                    )
                    .arg(
                        Arg::with_name("no-wait")
                            .long("no-wait")
                            .help("command-executer won't wait the command's result")
                            .takes_value(false),
                    ),
            )
            .subcommand(
                SubCommand::with_name("recv-file")
                    .about("Receive a file from the enclave")
                    .arg(
                        Arg::with_name("port")
                            .long("port")
                            .help("port")
                            .takes_value(true)
                            .required(true),
                    )
                    .arg(
                        Arg::with_name("cid")
                            .long("cid")
                            .help("cid")
                            .takes_value(true)
                            .required(true),
                    )
                    .arg(
                        Arg::with_name("localpath")
                            .long("localpath")
                            .help("localpath")
                            .takes_value(true)
                            .required(true),
                    )
                    .arg(
                        Arg::with_name("remotepath")
                            .long("remotepath")
                            .help("remotepath")
                            .takes_value(true)
                            .required(true),
                    ),
            )
            .subcommand(
                SubCommand::with_name("send-file")
                    .about("Send a file to the enclave")
                    .arg(
                        Arg::with_name("port")
                            .long("port")
                            .help("port")
                            .takes_value(true)
                            .required(true),
                    )
                    .arg(
                        Arg::with_name("cid")
                            .long("cid")
                            .help("cid")
                            .takes_value(true)
                            .required(true),
                    )
                    .arg(
                        Arg::with_name("localpath")
                            .long("localpath")
                            .help("localpath")
                            .takes_value(true)
                            .required(true),
                    )
                    .arg(
                        Arg::with_name("remotepath")
                            .long("remotepath")
                            .help("remotepath")
                            .takes_value(true)
                            .required(true),
                    ),
            )
    };
}


// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
#![deny(missing_docs)]
#![deny(warnings)]

use log::{debug, info};
use nix::sys::epoll;
use nix::sys::epoll::{EpollEvent, EpollFlags, EpollOp};
use nix::unistd::*;
use serde::de::DeserializeOwned;
use serde::Serialize;
use std::borrow::BorrowMut;
use std::fs;
use std::io::ErrorKind;
use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd};
use std::os::unix::net::UnixStream;

use crate::common::commands_parser::EmptyArgs;
use crate::common::logger::EnclaveProcLogWriter;
use crate::common::{
    enclave_proc_command_send_single, get_socket_path, get_sockets_dir_path, notify_error,
    read_u64_le, receive_from_stream,
};
use crate::common::{
    EnclaveProcessCommandType, EnclaveProcessReply, NitroCliErrorEnum, NitroCliFailure,
    NitroCliResult,
};
use crate::common::{ENCLAVE_PROC_WAIT_TIMEOUT_MSEC, MSG_ENCLAVE_CONFIRM};
use crate::enclave_proc::enclave_process_run;
use crate::new_nitro_cli_failure;

/// Spawn an enclave process and wait until it has detached and has
/// taken ownership of its communication socket.
pub fn enclave_proc_spawn(logger: &EnclaveProcLogWriter) -> NitroCliResult<UnixStream> {
    let (cli_socket, enclave_proc_socket) = UnixStream::pair().map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Could not create a socket pair: {:?}", e),
            NitroCliErrorEnum::SocketPairCreationFailure
        )
    })?;

    // Prevent the descriptor from being closed when calling exec().
    let enclave_proc_fd = enclave_proc_socket.as_raw_fd();
    unsafe {
        let flags = libc::fcntl(enclave_proc_fd, libc::F_GETFD);
        libc::fcntl(enclave_proc_fd, libc::F_SETFD, flags & !libc::FD_CLOEXEC);
    }

    // Spawn an intermediate child process. This will fork again in order to
    // create the detached enclave process.
    // Safety: enclave_proc_spawn is called early on and nitro-cli is not that this point
    // multi-threaded, which should prevent the issues around forking. However,
    // the safe way to do this would be to use a safe alternative such as Command to
    // re-execute this same process with another set of parameters.
    let fork_status = unsafe { fork() };

    if let Ok(ForkResult::Child) = fork_status {
        // This is our intermediate child process.
        enclave_process_run(enclave_proc_socket, logger);
    } else {
        fork_status.map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to create intermediate process: {:?}", e),
                NitroCliErrorEnum::ProcessSpawnFailure
            )
        })?;
    }

    // The enclave process will open a socket named "<enclave_id>.sock", but this
    // will only become available after the enclave has been successfully launched.
    // Until then, we can only use the pre-initialized socket pair to communicate
    // with the new process.
    Ok(cli_socket)
}

/// Connect to all existing enclave processes, returning a connection to each.
pub fn enclave_proc_connect_to_all() -> NitroCliResult<Vec<UnixStream>> {
    let paths = fs::read_dir(get_sockets_dir_path()).map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Failed to access sockets directory: {:?}", e),
            NitroCliErrorEnum::ReadFromDiskFailure
        )
    })?;
    Ok(paths
        .filter_map(|path| path.ok())
        .map(|path| path.path())
        .filter(|path| !path.is_dir())
        .filter_map(|path| {
            // Get the file path string.
            if let Some(path_str) = path.to_str() {
                // Enclave process sockets are named "<enclave_id>.sock".
                if !path_str.ends_with(".sock") {
                    return None;
                }

                // At this point we have found a potential socket.
                match UnixStream::connect(path_str) {
                    Ok(conn) => {
                        // We have connected to an enclave process.
                        info!("Connected to: {}", path_str);
                        return Some(conn);
                    }
                    Err(e) => {
                        if e.kind() == ErrorKind::PermissionDenied {
                            // Don't touch the socket if connection failed due to insufficient permissions.
                            info!("Connection to '{}' failed: {}", path_str, e);
                        } else {
                            // For all other connection errors, assume the socket is stale and delete it.
                            info!("Deleting stale socket: {}", path_str);
                            let _ = fs::remove_file(path_str).map_err(|e| {
                                new_nitro_cli_failure!(
                                    &format!("Failed to delete socket: {:?}", e),
                                    NitroCliErrorEnum::FileOperationFailure
                                )
                                .add_info(vec![path_str, "Remove"])
                            });
                        }
                    }
                }
            }

            None
        })
        .collect())
}

/// Open a connection to an enclave-specific socket.
pub fn enclave_proc_connect_to_single(enclave_id: &str) -> NitroCliResult<UnixStream> {
    let socket_path = get_socket_path(enclave_id).map_err(|e| {
        e.add_subaction("Connect to specific enclave process".to_string())
            .set_error_code(NitroCliErrorEnum::SocketError)
    })?;
    UnixStream::connect(socket_path).map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Failed to connect to specific enclave process: {:?}", e),
            NitroCliErrorEnum::SocketError
        )
    })
}

/// Broadcast a command to all available enclave processes.
pub fn enclave_proc_command_send_all<T>(
    cmd: EnclaveProcessCommandType,
    args: Option<&T>,
) -> NitroCliResult<(Vec<UnixStream>, usize)>
where
    T: Serialize,
{
    // Open a connection to each valid socket.
    let mut replies: Vec<UnixStream> = vec![];
    let epoll_fd = epoll::epoll_create().map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Failed to create epoll: {:?}", e),
            NitroCliErrorEnum::EpollError
        )
    })?;
    let comms: Vec<NitroCliResult<()>> = enclave_proc_connect_to_all()
        .map_err(|e| {
            e.add_subaction("Failed to send command to all enclave processes".to_string())
        })?
        .into_iter()
        .map(|mut socket| {
            // Send the command.
            enclave_proc_command_send_single(cmd, args, socket.borrow_mut())?;

            let raw_fd = socket.into_raw_fd();
            let mut process_evt = EpollEvent::new(EpollFlags::EPOLLIN, raw_fd as u64);

            // Add each valid connection to epoll.
            epoll::epoll_ctl(epoll_fd, EpollOp::EpollCtlAdd, raw_fd, &mut process_evt).map_err(
                |e| {
                    new_nitro_cli_failure!(
                        &format!("Failed to register socket with epoll: {:?}", e),
                        NitroCliErrorEnum::EpollError
                    )
                },
            )?;

            Ok(())
        })
        .collect();

    // Don't proceed unless at least one connection has been established.
    if comms.is_empty() {
        return Ok((vec![], 0));
    }

    // Get the number of transmission errors.
    let mut num_errors = comms.iter().filter(|result| result.is_err()).count();

    // Get the number of expected replies.
    let mut num_replies_expected = comms.len() - num_errors;
    let mut events = vec![EpollEvent::empty(); 1];

    while num_replies_expected > 0 {
        let num_events = loop {
            match epoll::epoll_wait(epoll_fd, &mut events[..], ENCLAVE_PROC_WAIT_TIMEOUT_MSEC) {
                Ok(num_events) => break num_events,
                Err(nix::errno::Errno::EINTR) => continue,
                // TODO: Handle bad descriptors (closed remote connections).
                Err(e) => {
                    return Err(new_nitro_cli_failure!(
                        &format!("Failed to wait on epoll: {:?}", e),
                        NitroCliErrorEnum::EpollError
                    ))
                }
            }
        };

        // We will handle this reply, irrespective of its status (successful or failed).
        num_replies_expected -= 1;

        // Check if a time-out has occurred.
        if num_events == 0 {
            continue;
        }

        let input_stream_raw_fd = events[0].data() as RawFd;
        let mut input_stream = unsafe { UnixStream::from_raw_fd(input_stream_raw_fd) };

        // Handle the reply we received.
        if let Ok(reply) = read_u64_le(&mut input_stream) {
            if reply == MSG_ENCLAVE_CONFIRM {
                debug!("Got confirmation from {:?}", input_stream);
                replies.push(input_stream);
            }
        }

        epoll::epoll_ctl(
            epoll_fd,
            EpollOp::EpollCtlDel,
            input_stream_raw_fd,
            Option::None,
        )
        .map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to remove socket from epoll: {:?}", e),
                NitroCliErrorEnum::EpollError
            )
        })?;
    }

    // Update the number of connections that have yielded errors.
    num_errors = comms.len() - replies.len();

    Ok((replies, num_errors))
}

/// Print the output from a single enclave process.
fn enclave_proc_handle_output<T>(conn: &mut UnixStream) -> (Option<T>, Option<i32>)
where
    T: DeserializeOwned,
{
    let mut stdout_str = String::new();
    let mut status: Option<i32> = None;

    // The contents meant for standard output must always form a valid JSON object.
    while let Ok(reply) = receive_from_stream::<EnclaveProcessReply>(conn) {
        match reply {
            EnclaveProcessReply::StdOutMessage(msg) => stdout_str.push_str(&msg),
            EnclaveProcessReply::StdErrMessage(msg) => eprint!("{}", msg),
            EnclaveProcessReply::Status(status_code) => status = Some(status_code),
        }
    }

    // Shut the connection down.
    match conn.shutdown(std::net::Shutdown::Both) {
        Ok(()) => (),
        Err(e) => {
            notify_error(&format!("Failed to shut connection down: {}", e));
            status = Some(-1);
        }
    }

    // Decode the JSON object.
    let json_obj = serde_json::from_str::<T>(&stdout_str).ok();
    (json_obj, status)
}

/// Fetch JSON objects and statuses from all connected enclave processes.
pub fn enclave_proc_handle_outputs<T>(conns: &mut [UnixStream]) -> Vec<(T, i32)>
where
    T: DeserializeOwned,
{
    let mut objects: Vec<(T, i32)> = Vec::new();

    for conn in conns.iter_mut() {
        // We only count connections that have yielded a valid JSON object and a status
        let (object, status) = enclave_proc_handle_output::<T>(conn);
        if let Some(object) = object {
            if let Some(status) = status {
                objects.push((object, status));
            }
        }
    }

    objects
}

/// Process reply messages from all connected enclave processes.
pub fn enclave_process_handle_all_replies<T>(
    replies: &mut [UnixStream],
    prev_failed_conns: usize,
    print_as_vec: bool,
    allowed_return_codes: Vec<i32>,
) -> NitroCliResult<Vec<T>>
where
    T: Clone + DeserializeOwned + Serialize,
{
    let objects = enclave_proc_handle_outputs::<T>(replies);
    let failed_conns = prev_failed_conns + replies.len() - objects.len();

    // Print a message if we have any connections that have failed.
    if failed_conns > 0 {
        eprintln!("Failed connections: {}", failed_conns);
    }

    // Output the received objects either individually or as an array.
    if print_as_vec {
        let obj_vec: Vec<T> = objects.iter().map(|v| v.0.clone()).collect();
        println!(
            "{}",
            serde_json::to_string_pretty(&obj_vec).map_err(|e| new_nitro_cli_failure!(
                &format!("Failed to print JSON vector: {:?}", e),
                NitroCliErrorEnum::SerdeError
            ))?
        );
    } else {
        for object in objects.iter().map(|v| v.0.clone()) {
            println!(
                "{}",
                serde_json::to_string_pretty(&object).map_err(|e| new_nitro_cli_failure!(
                    &format!("Failed to print JSON object: {:?}", e),
                    NitroCliErrorEnum::SerdeError
                ))?
            );
        }
    }

    // We fail on any error codes or failed connections.
    if objects
        .iter()
        .filter(|v| !allowed_return_codes.contains(&v.1))
        .count()
        > 0
    {
        return Err(new_nitro_cli_failure!(
            &format!(
                "Failed to execute {} enclave process commands",
                objects
                    .iter()
                    .filter(|v| !allowed_return_codes.contains(&v.1))
                    .count()
            ),
            NitroCliErrorEnum::EnclaveProcessCommandNotExecuted
        ));
    } else if failed_conns > 0 {
        return Err(new_nitro_cli_failure!(
            &format!("Failed to connect to {} enclave processes", failed_conns),
            NitroCliErrorEnum::EnclaveProcessConnectionFailure
        ));
    }

    Ok(objects.into_iter().map(|(o, _)| o).collect())
}

/// Obtain an enclave's CID given its full ID.
pub fn enclave_proc_get_cid(enclave_id: &str) -> NitroCliResult<u64> {
    let mut comm = enclave_proc_connect_to_single(enclave_id)
        .map_err(|e| e.add_subaction("Failed to connect to enclave process".to_string()))?;
    // TODO: Replicate output of old CLI on invalid enclave IDs.
    enclave_proc_command_send_single::<EmptyArgs>(
        EnclaveProcessCommandType::GetEnclaveCID,
        None,
        &mut comm,
    )
    .map_err(|e| e.add_subaction("Failed to send CID request to enclave process".to_string()))?;

    info!("Sent command: GetEnclaveCID");
    let enclave_cid = read_u64_le(&mut comm)
        .map_err(|e| e.add_subaction(String::from("Failed to read CID from enclave process")))?;

    // We got the CID, so shut the connection down.
    comm.shutdown(std::net::Shutdown::Both).map_err(|e| {
        new_nitro_cli_failure!(
            &format!(
                "Failed to shut down connection after obtaining enclave CID: {:?}",
                e
            ),
            NitroCliErrorEnum::SocketError
        )
    })?;

    Ok(enclave_cid)
}

/// Obtain an enclave's flags given its full ID.
pub fn enclave_proc_get_flags(enclave_id: &str) -> NitroCliResult<u64> {
    let mut comm = enclave_proc_connect_to_single(enclave_id)
        .map_err(|e| e.add_subaction("Failed to connect to enclave process".to_string()))?;
    // TODO: Replicate output of old CLI on invalid enclave IDs.
    enclave_proc_command_send_single::<EmptyArgs>(
        EnclaveProcessCommandType::GetEnclaveFlags,
        None,
        &mut comm,
    )
    .map_err(|e| e.add_subaction("Failed to send flags request to enclave process".to_string()))?;

    info!("Sent command: GetEnclaveFlags");
    let enclave_flags = read_u64_le(&mut comm)
        .map_err(|e| e.add_subaction(String::from("Failed to read flags from enclave process")))?;

    // We got the flags, so shut the connection down.
    comm.shutdown(std::net::Shutdown::Both).map_err(|e| {
        new_nitro_cli_failure!(
            &format!(
                "Failed to shut down connection after obtaining enclave flags: {:?}",
                e
            ),
            NitroCliErrorEnum::SocketError
        )
    })?;

    Ok(enclave_flags)
}


// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
#![deny(missing_docs)]
#![deny(warnings)]
#![allow(clippy::too_many_arguments)]

//! This crate provides the functionality for the Nitro CLI process.

/// The common module (shared between the CLI and enclave process).
pub mod common;
/// The enclave process module.
pub mod enclave_proc;
/// The module covering the communication between a CLI instance and enclave processes.
pub mod enclave_proc_comm;
/// The CLI-specific utilities module.
pub mod utils;

use aws_nitro_enclaves_image_format::defs::eif_hasher::EifHasher;
use aws_nitro_enclaves_image_format::utils::eif_reader::EifReader;
use aws_nitro_enclaves_image_format::{generate_build_info, utils::get_pcrs};
use log::{debug, info};
use sha2::{Digest, Sha384};
use std::collections::BTreeMap;
use std::convert::TryFrom;
use std::fs::{File, OpenOptions};
use std::io::{self, Read, Write};
use std::os::unix::net::UnixStream;
use std::path::PathBuf;

use common::commands_parser::{BuildEnclavesArgs, EmptyArgs, RunEnclavesArgs};
use common::json_output::{
    EifDescribeInfo, EnclaveBuildInfo, EnclaveTerminateInfo, MetadataDescribeInfo,
};
use common::{enclave_proc_command_send_single, get_sockets_dir_path};
use common::{EnclaveProcessCommandType, NitroCliErrorEnum, NitroCliFailure, NitroCliResult};
use enclave_proc_comm::{
    enclave_proc_command_send_all, enclave_proc_handle_outputs, enclave_process_handle_all_replies,
};

use utils::{Console, PcrType};

/// Hypervisor CID as defined by <http://man7.org/linux/man-pages/man7/vsock.7.html>.
pub const VMADDR_CID_HYPERVISOR: u32 = 0;

/// An offset applied to an enclave's CID in order to determine its console port.
pub const CID_TO_CONSOLE_PORT_OFFSET: u32 = 10000;

/// Default blobs path to be used if the corresponding environment variable is not set.
const DEFAULT_BLOBS_PATH: &str = "/usr/share/nitro_enclaves/blobs/";

/// Build an enclave image file with the provided arguments.
pub fn build_enclaves(args: BuildEnclavesArgs) -> NitroCliResult<()> {
    debug!("build_enclaves");
    eprintln!("Start building the Enclave Image...");
    build_from_docker(
        &args.docker_uri,
        &args.docker_dir,
        &args.output,
        &args.signing_certificate,
        &args.private_key,
        &args.img_name,
        &args.img_version,
        &args.metadata,
    )
    .map_err(|e| e.add_subaction("Failed to build EIF from docker".to_string()))?;
    Ok(())
}

/// Build an enclave image file from a Docker image.
pub fn build_from_docker(
    docker_uri: &str,
    docker_dir: &Option<String>,
    output_path: &str,
    signing_certificate: &Option<String>,
    private_key: &Option<String>,
    img_name: &Option<String>,
    img_version: &Option<String>,
    metadata_path: &Option<String>,
) -> NitroCliResult<(File, BTreeMap<String, String>)> {
    let blobs_path =
        blobs_path().map_err(|e| e.add_subaction("Failed to retrieve blobs path".to_string()))?;
    let cmdline_file_path = format!("{}/cmdline", blobs_path);
    let mut cmdline_file = File::open(cmdline_file_path.clone()).map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Could not open kernel command line file: {:?}", e),
            NitroCliErrorEnum::FileOperationFailure
        )
        .add_info(vec![&cmdline_file_path, "Open"])
    })?;

    let mut cmdline = String::new();
    cmdline_file.read_to_string(&mut cmdline).map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Failed to read kernel command line: {:?}", e),
            NitroCliErrorEnum::FileOperationFailure
        )
        .add_info(vec![&cmdline_file_path, "Read"])
    })?;

    let mut file_output = OpenOptions::new()
        .read(true)
        .create(true)
        .write(true)
        .truncate(true)
        .open(output_path)
        .map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Could not create output file: {:?}", e),
                NitroCliErrorEnum::FileOperationFailure
            )
            .add_info(vec![output_path, "Open"])
        })?;

    let kernel_image_name = match std::env::consts::ARCH {
        "aarch64" => "Image",
        "x86_64" => "bzImage",
        _ => "undefined",
    };

    let kernel_path = format!("{}/{}", blobs_path, kernel_image_name);
    let build_info = generate_build_info!(&format!("{}.config", kernel_path)).map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Could not generate build info: {:?}", e),
            NitroCliErrorEnum::EifBuildingError
        )
    })?;

    let mut docker2eif = enclave_build::Docker2Eif::new(
        docker_uri.to_string(),
        format!("{}/init", blobs_path),
        format!("{}/nsm.ko", blobs_path),
        kernel_path,
        cmdline.trim().to_string(),
        format!("{}/linuxkit", blobs_path),
        &mut file_output,
        artifacts_path()?,
        signing_certificate,
        private_key,
        img_name.clone(),
        img_version.clone(),
        metadata_path.clone(),
        build_info,
    )
    .map_err(|err| {
        new_nitro_cli_failure!(
            &format!("Failed to create EIF image: {:?}", err),
            NitroCliErrorEnum::EifBuildingError
        )
    })?;

    if let Some(docker_dir) = docker_dir {
        docker2eif
            .build_docker_image(docker_dir.clone())
            .map_err(|err| {
                new_nitro_cli_failure!(
                    &format!("Failed to build docker image: {:?}", err),
                    NitroCliErrorEnum::DockerImageBuildError
                )
            })?;
    } else {
        docker2eif.pull_docker_image().map_err(|err| {
            new_nitro_cli_failure!(
                &format!("Failed to pull docker image: {:?}", err),
                NitroCliErrorEnum::DockerImagePullError
            )
        })?;
    }
    let measurements = docker2eif.create().map_err(|err| {
        new_nitro_cli_failure!(
            &format!("Failed to create EIF image: {:?}", err),
            NitroCliErrorEnum::EifBuildingError
        )
    })?;
    eprintln!("Enclave Image successfully created.");

    let info = EnclaveBuildInfo::new(measurements.clone());
    println!(
        "{}",
        serde_json::to_string_pretty(&info).map_err(|err| new_nitro_cli_failure!(
            &format!("Failed to display EnclaveBuild data: {:?}", err),
            NitroCliErrorEnum::SerdeError
        ))?
    );

    Ok((file_output, measurements))
}

/// Creates new enclave name
///
/// Requests the names of all running instances and checks the
/// occurrence of the chosen name for the new enclave.
pub fn new_enclave_name(run_args: RunEnclavesArgs, names: Vec<String>) -> NitroCliResult<String> {
    let enclave_name = match run_args.enclave_name {
        Some(enclave_name) => enclave_name,
        None => {
            // Get name of EIF file from path eg. path/to/eif/hello.eif -> hello
            // If the extension is missing, the whole file name will be chosen
            let path_split: Vec<&str> = run_args.eif_path.split('/').collect();
            path_split[path_split.len() - 1]
                .trim_end_matches(".eif")
                .to_string()
        }
    };

    let mut idx = 0;
    let mut result_name = enclave_name.clone();

    // If duplicates are found, add index to name eg. testName -> testName_1 -> testName_2 ..
    while names.contains(&result_name) {
        idx += 1;
        result_name = enclave_name.clone() + &'_'.to_string() + &idx.to_string();
    }

    Ok(result_name)
}

/// Returns information related to the given EIF
///
/// Calculates PCRs 0, 1, 2, 8 at each call in addition to metadata,
/// EIF details, identification provided by the user at build.
pub fn describe_eif(eif_path: String) -> NitroCliResult<EifDescribeInfo> {
    let mut eif_reader = EifReader::from_eif(eif_path).map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Failed to initialize EIF reader: {:?}", e),
            NitroCliErrorEnum::EifParsingError
        )
    })?;
    let measurements = get_pcrs(
        &mut eif_reader.image_hasher,
        &mut eif_reader.bootstrap_hasher,
        &mut eif_reader.app_hasher,
        &mut eif_reader.cert_hasher,
        Sha384::new(),
        eif_reader.signature_section.is_some(),
    )
    .map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Failed to get PCR values: {:?}", e),
            NitroCliErrorEnum::EifParsingError
        )
    })?;

    let mut describe_meta: Option<MetadataDescribeInfo> = None;
    let mut img_name: Option<String> = None;
    let mut img_version: Option<String> = None;

    if let Some(meta) = eif_reader.get_metadata() {
        img_name = Some(meta.img_name.clone());
        img_version = Some(meta.img_version.clone());
        describe_meta = Some(MetadataDescribeInfo::new(meta));
    }

    let mut info = EifDescribeInfo {
        version: eif_reader.get_header().version,
        build_info: EnclaveBuildInfo::new(measurements.clone()),
        is_signed: false,
        cert_info: None,
        crc_check: eif_reader.check_crc(),
        sign_check: None,
        img_name,
        img_version,
        metadata: describe_meta,
    };

    // Check if signature section is present
    if measurements.get(&"PCR8".to_string()).is_some() {
        let cert_info = eif_reader
            .get_certificate_info(measurements)
            .map_err(|err| {
                new_nitro_cli_failure!(
                    &format!("Failed to get certificate sigining info: {:?}", err),
                    NitroCliErrorEnum::EifParsingError
                )
            })?;
        info.is_signed = true;
        info.cert_info = Some(cert_info);
        info.sign_check = eif_reader.sign_check;
    }

    println!(
        "{}",
        serde_json::to_string_pretty(&info)
            .map_err(|err| {
                new_nitro_cli_failure!(
                    &format!("Failed to display EIF describe data: {:?}", err),
                    NitroCliErrorEnum::SerdeError
                )
            })?
            .as_str(),
    );

    Ok(info)
}

/// Returns the value of the `NITRO_CLI_BLOBS` environment variable.
///
/// This variable specifies where all the blobs necessary for building
/// an enclave image are stored. As of now the blobs are:
/// - *bzImage*: A kernel image if the local arch is x86_64 or
/// - *Image*  : A kernel image if the local arch is aarch64
/// - *init*: The initial init process that is bootstraping the environment.
/// - *linuxkit*: A slightly modified version of linuxkit.
/// - *cmdline*: A file containing the kernel commandline.
fn blobs_path() -> NitroCliResult<String> {
    // TODO Improve error message with a suggestion to the user
    // consider using the default path used by rpm install
    let blobs_res = std::env::var("NITRO_CLI_BLOBS");

    Ok(blobs_res.unwrap_or_else(|_| DEFAULT_BLOBS_PATH.to_string()))
}

/// Returns the value of the `NITRO_CLI_ARTIFACTS` environment variable.
///
/// This variable configures the path where the build artifacts should be saved.
fn artifacts_path() -> NitroCliResult<String> {
    if let Ok(artifacts) = std::env::var("NITRO_CLI_ARTIFACTS") {
        std::fs::create_dir_all(artifacts.clone()).map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Could not create artifacts path {}: {:?}", artifacts, e),
                NitroCliErrorEnum::FileOperationFailure
            )
            .add_info(vec![&artifacts, "Create"])
        })?;
        Ok(artifacts)
    } else if let Ok(home) = std::env::var("HOME") {
        let artifacts = format!("{}/.nitro_cli/", home);
        std::fs::create_dir_all(artifacts.clone()).map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Could not create artifacts path {}: {:?}", artifacts, e),
                NitroCliErrorEnum::FileOperationFailure
            )
            .add_info(vec![&artifacts, "Create"])
        })?;
        Ok(artifacts)
    } else {
        Err(new_nitro_cli_failure!(
            "Could not find a folder for the CLI artifacts, set either HOME or NITRO_CLI_ARTIFACTS",
            NitroCliErrorEnum::ArtifactsPathNotSet
        ))
    }
}

/// Wrapper over the console connection function.
pub fn console_enclaves(
    enclave_cid: u64,
    disconnect_timeout_sec: Option<u64>,
) -> NitroCliResult<()> {
    debug!("console_enclaves");
    println!("Connecting to the console for enclave {}...", enclave_cid);
    enclave_console(enclave_cid, disconnect_timeout_sec)?;
    Ok(())
}

/// Connects to the enclave console and prints it continously.
pub fn enclave_console(
    enclave_cid: u64,
    disconnect_timeout_sec: Option<u64>,
) -> NitroCliResult<()> {
    let console = Console::new(
        VMADDR_CID_HYPERVISOR,
        u32::try_from(enclave_cid).map_err(|err| {
            new_nitro_cli_failure!(
                &format!("Failed to parse enclave CID: {:?}", err),
                NitroCliErrorEnum::IntegerParsingError
            )
        })? + CID_TO_CONSOLE_PORT_OFFSET,
    )
    .map_err(|e| e.add_subaction("Connect to enclave console".to_string()))?;
    println!("Successfully connected to the console.");
    console
        .read_to(io::stdout().by_ref(), disconnect_timeout_sec)
        .map_err(|e| e.add_subaction("Connect to enclave console".to_string()))?;

    Ok(())
}

/// Terminates all enclave instances belonging to the current user (or all
/// instances, if the current user has `root` permissions).
pub fn terminate_all_enclaves() -> NitroCliResult<()> {
    let sockets_dir = get_sockets_dir_path();
    let mut replies: Vec<UnixStream> = vec![];
    let sockets = std::fs::read_dir(sockets_dir.as_path()).map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Error while accessing sockets directory: {:?}", e),
            NitroCliErrorEnum::FileOperationFailure
        )
        .add_info(vec![
            sockets_dir
                .as_path()
                .to_str()
                .unwrap_or("Invalid unicode directory name"),
            "Read",
        ])
    })?;

    let mut err_socket_files: usize = 0;
    let mut failed_connections: Vec<PathBuf> = Vec::new();
    for socket in sockets {
        let entry = match socket {
            Ok(value) => value,
            Err(_) => {
                err_socket_files += 1;
                continue;
            }
        };

        // Send a `terminate-enclave` command through each socket,
        // irrespective of the enclave process owner. The security policy
        // inside the enclave process is responsible with checking the
        // command's permissions.
        let mut stream = match UnixStream::connect(entry.path()) {
            Ok(value) => value,
            Err(_) => {
                failed_connections.push(entry.path());
                continue;
            }
        };

        if enclave_proc_command_send_single::<EmptyArgs>(
            EnclaveProcessCommandType::Terminate,
            None,
            &mut stream,
        )
        .is_err()
        {
            failed_connections.push(entry.path());
        } else {
            replies.push(stream);
        }
    }

    // Remove stale socket files.
    for stale_socket in &failed_connections {
        info!("Deleting stale socket: {:?}", stale_socket);
        let _ = std::fs::remove_file(stale_socket);
    }

    enclave_process_handle_all_replies::<EnclaveTerminateInfo>(
        &mut replies,
        failed_connections.len() + err_socket_files,
        false,
        vec![0, libc::EACCES],
    )
    .map_err(|e| e.add_subaction("Failed to handle all enclave processes replies".to_string()))
    .map(|_| ())
}

/// Queries all enclaves for their name
pub fn get_all_enclave_names() -> NitroCliResult<Vec<String>> {
    let (comms, _) =
        enclave_proc_command_send_all::<EmptyArgs>(EnclaveProcessCommandType::GetEnclaveName, None)
            .map_err(|e| {
                e.add_subaction(
                    "Failed to send GetEnclaveName command to all enclave processes".to_string(),
                )
                .set_action("Get Enclave Names".to_string())
            })?;

    let mut replies: Vec<UnixStream> = vec![];
    replies.extend(comms);
    let objects = enclave_proc_handle_outputs::<String>(&mut replies)
        .iter()
        .map(|v| v.0.clone())
        .collect();
    Ok(objects)
}

/// Sends the name to all the running enclaves and expects a response
/// with the ID of the one that uniquely matched
pub fn get_id_by_name(name: String) -> NitroCliResult<String> {
    let (comms, _) = enclave_proc_command_send_all::<String>(
        EnclaveProcessCommandType::GetIDbyName,
        Some(&name),
    )
    .map_err(|e| {
        e.add_subaction("Failed to send GetIDbyName command to all enclave processes".to_string())
            .set_action("Get Enclave Names".to_string())
    })?;

    let mut replies: Vec<UnixStream> = vec![];
    replies.extend(comms);
    let mut objects: Vec<String> = enclave_proc_handle_outputs::<String>(&mut replies)
        .iter()
        .map(|v| v.0.clone())
        .collect();

    // Check if the name was not found or if there are multiple matches
    if objects.len() != 1 {
        return Err(new_nitro_cli_failure!(
            match objects.len() {
                0 => "No enclave matched the given name.".to_string(),
                _ => "Conflicting enclave names have been found.".to_string(),
            },
            NitroCliErrorEnum::EnclaveNamingError
        ));
    }

    Ok(objects.remove(0))
}

/// For the given file, return the PCR value
///
/// Based on the pcr_type, calculate the PCR hash of the input. The default
/// type takes the bytes of the input file and adds them to the hasher.
/// The certificate type performs additional serialization before hashing.
pub fn get_file_pcr(path: String, pcr_type: PcrType) -> NitroCliResult<BTreeMap<String, String>> {
    let mut key = "PCR".to_string();
    // Initialize hasher
    let mut hasher = EifHasher::new_without_cache(Sha384::new()).map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Could not create hasher: {:?}", e),
            NitroCliErrorEnum::HasherError
        )
    })?;
    let mut file = File::open(path).map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Failed to open file: {:?}", e),
            NitroCliErrorEnum::FileOperationFailure
        )
    })?;
    let mut buf = Vec::new();
    file.read_to_end(&mut buf).map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Failed to read file: {:?}", e),
            NitroCliErrorEnum::FileOperationFailure
        )
    })?;
    // Treat the input buffer by PCR type
    match pcr_type {
        PcrType::DefaultType => {}
        PcrType::SigningCertificate => {
            key = "PCR8".to_string();
            let cert = openssl::x509::X509::from_pem(&buf[..]).map_err(|e| {
                new_nitro_cli_failure!(
                    &format!("Failed to deserialize .pem: {:?}", e),
                    NitroCliErrorEnum::HasherError
                )
            })?;
            buf = cert.to_der().map_err(|e| {
                new_nitro_cli_failure!(
                    &format!("Failed to serialize certificate: {:?}", e),
                    NitroCliErrorEnum::HasherError
                )
            })?;
        }
    }
    hasher.write_all(&buf).map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Could not write to hasher: {:?}", e),
            NitroCliErrorEnum::HasherError
        )
    })?;
    let hash = hex::encode(hasher.tpm_extend_finalize_reset().map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Could not get result for hasher: {:?}", e),
            NitroCliErrorEnum::HasherError
        )
    })?);

    let mut result = BTreeMap::new();
    result.insert(key, hash);
    println!(
        "{}",
        serde_json::to_string_pretty(&result)
            .map_err(|err| {
                new_nitro_cli_failure!(
                    &format!("Failed to display PCR(s): {:?}", err),
                    NitroCliErrorEnum::SerdeError
                )
            })?
            .as_str(),
    );
    Ok(result)
}

/// Macro defining the arguments configuration for a *Nitro CLI* application.
#[macro_export]
macro_rules! create_app {
    () => {
        App::new("Nitro CLI")
            .about("CLI for enclave lifetime management")
            .setting(AppSettings::ArgRequiredElseHelp)
            .version(env!("CARGO_PKG_VERSION"))
            .subcommand(
                SubCommand::with_name("run-enclave")
                    .about("Starts a new enclave")
                    .arg(
                        Arg::with_name("cpu-ids")
                            .long("cpu-ids")
                            .help("List of cpu-ids that will be provided to the enclave")
                            .takes_value(true)
                            .multiple(true)
                            .min_values(1)
                            .required_unless("cpu-count")
                            .required_unless("config")
                            .conflicts_with("cpu-count")
                            .conflicts_with("config"),
                    )
                    .arg(
                        Arg::with_name("cpu-count")
                            .long("cpu-count")
                            .help("Number of cpus")
                            .takes_value(true)
                            .required_unless("cpu-ids")
                            .required_unless("config")
                            .conflicts_with("cpu-ids")
                            .conflicts_with("config"),
                    )
                    .arg(
                        Arg::with_name("memory")
                            .long("memory")
                            .help(
                                "Memory to allocate for the enclave in MB. Depending on the available \
                                pages, more might be allocated."
                            )
                            .required_unless("config")
                            .takes_value(true)
                            .conflicts_with("config"),
                    )
                    .arg(
                        Arg::with_name("eif-path")
                            .long("eif-path")
                            .help("Path pointing to a prebuilt Eif image")
                            .required_unless("config")
                            .takes_value(true)
                            .conflicts_with("config"),
                    )
                    .arg(
                        Arg::with_name("enclave-cid")
                            .long("enclave-cid")
                            .takes_value(true)
                            .help("CID to be used for the newly started enclave")
                            .conflicts_with("config"),
                    )
                    .arg(
                        Arg::with_name("debug-mode")
                            .long("debug-mode")
                            .takes_value(false)
                            .help(
                                "Starts enclave in debug-mode. This makes the console of the enclave \
                                available over vsock at CID: VMADDR_CID_HYPERVISOR (0), port: \
                                enclave_cid + 10000. \n The stream could be accessed with the console \
                                sub-command"
                            )
                            .conflicts_with("config"),
                    )
                    .arg(
                        Arg::with_name("attach-console")
                            .long("attach-console")
                            .takes_value(false)
                            .help(
                                "Attach the enclave console immediately after starting the enclave. \
                                (implies debug-mode)"
                            )
                    )
                    .arg(
                        Arg::with_name("enclave-name")
                            .long("enclave-name")
                            .takes_value(true)
                            .help("Custom name assigned to the enclave by the user")
                            .required(false)
                            .conflicts_with("config"),
                    )
                    .arg(
                        Arg::with_name("config")
                            .long("config")
                            .takes_value(true)
                            .value_name("json-config")
                            .help("Config is used to read enclave settings from JSON file"),
                    ),
            )
            .subcommand(
                SubCommand::with_name("terminate-enclave")
                    .about("Terminates an enclave")
                    .arg(
                        Arg::with_name("enclave-id")
                            .long("enclave-id")
                            .takes_value(true)
                            .help("Enclave ID, used to uniquely identify an enclave")
                            .required_unless("all")
                            .required_unless("enclave-name")
                            .conflicts_with("all")
                            .conflicts_with("enclave-name"),
                    )
                    .arg(
                        Arg::with_name("all")
                            .long("all")
                            .takes_value(false)
                            .help("Terminate all running enclave instances belonging to the current user")
                            .required_unless("enclave-id")
                            .required_unless("enclave-name")
                            .conflicts_with("enclave-id")
                            .conflicts_with("enclave-name"),
                    )
                    .arg(
                        Arg::with_name("enclave-name")
                            .long("enclave-name")
                            .takes_value(true)
                            .help("Enclave name, used to uniquely identify an enclave")
                            .required_unless("enclave-id")
                            .required_unless("all")
                            .conflicts_with("enclave-id")
                            .conflicts_with("all"),
                    ),
            )
            .subcommand(
                SubCommand::with_name("build-enclave")
                    .about("Builds an enclave image and saves it to a file")
                    .arg(
                        Arg::with_name("docker-uri")
                            .long("docker-uri")
                            .help(
                                "Uri pointing to an existing docker container or to be created  \
                                 locally when docker-dir is present",
                            )
                            .required(true)
                            .takes_value(true),
                    )
                    .arg(
                        Arg::with_name("docker-dir")
                            .long("docker-dir")
                            .help("Local path to a directory containing a Dockerfile")
                            .takes_value(true),
                    )
                    .arg(
                        Arg::with_name("output-file")
                            .long("output-file")
                            .help("Location where the Enclave Image should be saved")
                            .group("action")
                            .required(true)
                            .takes_value(true),
                    )
                    .arg(
                        Arg::with_name("signing-certificate")
                            .long("signing-certificate")
                            .help("Local path to developer's X509 signing certificate.")
                            .takes_value(true),
                    )
                    .arg(
                        Arg::with_name("private-key")
                            .long("private-key")
                            .help("Local path to developer's Eliptic Curve private key.")
                            .takes_value(true),
                    )
                    .arg(
                        Arg::with_name("image_name")
                            .long("name")
                            .help("Name for enclave image")
                            .takes_value(true),
                    )
                    .arg(
                        Arg::with_name("image_version")
                            .long("version")
                            .help("Version of the enclave image")
                            .takes_value(true),
                    )
                    .arg(
                        Arg::with_name("metadata")
                            .long("metadata")
                            .help("Path to JSON containing the custom metadata provided by the user.")
                            .takes_value(true),
                    ),
            )
            .subcommand(
                SubCommand::with_name("describe-eif")
                    .about("Returns information about the EIF found at a given path.")
                    .arg(
                        Arg::with_name("eif-path")
                            .long("eif-path")
                            .help("Path to the EIF to describe.")
                            .required(true)
                            .takes_value(true),
                    ),
            )
            .subcommand(
                SubCommand::with_name("describe-enclaves")
                    .about("Returns a list of the running enclaves")
                    .arg(
                        Arg::with_name("metadata")
                            .long("metadata")
                            .help("Adds EIF metadata of the current enclaves to the command output.")
                        ),
            )
            .subcommand(
                SubCommand::with_name("console")
                    .about("Connect to the console of an enclave")
                    .arg(
                        Arg::with_name("enclave-id")
                            .long("enclave-id")
                            .takes_value(true)
                            .help("Enclave ID, used to uniquely identify an enclave")
                            .required_unless("enclave-name")
                            .conflicts_with("enclave-name"),
                    )
                    .arg(
                        Arg::with_name("disconnect-timeout")
                            .long("disconnect-timeout")
                            .takes_value(true)
                            .help("The time in seconds after the console disconnects from the enclave")
                            .required(false)
                    )
                    .arg(
                        Arg::with_name("enclave-name")
                            .long("enclave-name")
                            .takes_value(true)
                            .help("Enclave name, used to uniquely identify an enclave")
                            .required_unless("enclave-id")
                            .conflicts_with("enclave-id"),
                    ),
            )
            .subcommand(
                SubCommand::with_name("pcr")
                    .about("Return the PCR hash value of the given input")
                    .arg(
                        Arg::with_name("signing-certificate")
                            .long("signing-certificate")
                            .takes_value(true)
                            .help("Takes the path to the '.pem' signing certificate and returns PCR8. Can be used to identify the certificate used to sign an EIF")
                            .required_unless("input")
                            .conflicts_with("input"),
                    )
                    .arg(
                        Arg::with_name("input")
                            .long("input")
                            .takes_value(true)
                            .help("Given a path to a file, returns the PCR hash of the bytes it contains")
                            .required_unless("signing-certificate")
                            .conflicts_with("signing-certificate"),
                    ),
            )
            .subcommand(
                SubCommand::with_name("explain")
                    .about("Display detailed information about an error returned by a misbehaving Nitro CLI command")
                    .arg(
                        Arg::with_name("error-code")
                            .long("error-code")
                            .takes_value(true)
                            .help("Error code, as returned by the misbehaving Nitro CLI command")
                            .required(true),
                    ),
            )
    };
}


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

#![deny(missing_docs)]
#![deny(warnings)]

//! This is the entry point for the Nitro CLI process.

extern crate lazy_static;

use clap::{App, AppSettings, Arg, SubCommand};
use log::info;
use std::os::unix::net::UnixStream;

use nitro_cli::common::commands_parser::{
    BuildEnclavesArgs, ConsoleArgs, DescribeEnclavesArgs, EmptyArgs, ExplainArgs, PcrArgs,
    RunEnclavesArgs, TerminateEnclavesArgs,
};
use nitro_cli::common::document_errors::explain_error;
use nitro_cli::common::json_output::{EnclaveDescribeInfo, EnclaveRunInfo, EnclaveTerminateInfo};
use nitro_cli::common::{
    enclave_proc_command_send_single, logger, NitroCliErrorEnum, NitroCliFailure, NitroCliResult,
};
use nitro_cli::common::{EnclaveProcessCommandType, ExitGracefully};
use nitro_cli::enclave_proc::resource_manager::NE_ENCLAVE_DEBUG_MODE;
use nitro_cli::enclave_proc_comm::{
    enclave_proc_command_send_all, enclave_proc_connect_to_single, enclave_proc_get_cid,
    enclave_proc_get_flags, enclave_proc_spawn, enclave_process_handle_all_replies,
};
use nitro_cli::{
    build_enclaves, console_enclaves, create_app, describe_eif, get_all_enclave_names,
    get_file_pcr, new_enclave_name, new_nitro_cli_failure, terminate_all_enclaves,
};

const RUN_ENCLAVE_STR: &str = "Run Enclave";
const DESCRIBE_ENCLAVE_STR: &str = "Describe Enclave";
const DESCRIBE_EIF_STR: &str = "Describe EIF";
const TERMINATE_ENCLAVE_STR: &str = "Terminate Enclave";
const TERMINATE_ALL_ENCLAVES_STR: &str = "Terminate All Enclaves";
const BUILD_ENCLAVE_STR: &str = "Build Enclave";
const ENCLAVE_CONSOLE_STR: &str = "Enclave Console";
const EXPLAIN_ERR_STR: &str = "Explain Error";
const NEW_NAME_STR: &str = "New Enclave Name";
const FILE_PCR_STR: &str = "File PCR";

/// *Nitro CLI* application entry point.
fn main() {
    let version_str: String = env!("CARGO_PKG_VERSION").to_string();

    // Command-line specification for the Nitro CLI.
    let mut app = create_app!();
    app = app.version(&*version_str);
    let args = app.get_matches();
    let logger = logger::init_logger()
        .map_err(|e| e.set_action("Logger initialization".to_string()))
        .ok_or_exit_with_errno(None);
    let mut replies: Vec<UnixStream> = vec![];

    logger
        .update_logger_id(format!("nitro-cli:{}", std::process::id()).as_str())
        .map_err(|e| e.set_action("Update CLI Process Logger ID".to_string()))
        .ok_or_exit_with_errno(None);
    info!("Start Nitro CLI");

    match args.subcommand() {
        Some(("run-enclave", args)) => {
            let mut run_args = RunEnclavesArgs::new_with(args)
                .map_err(|err| {
                    err.add_subaction("Failed to construct RunEnclave arguments".to_string())
                        .set_action(RUN_ENCLAVE_STR.to_string())
                })
                .ok_or_exit_with_errno(None);
            let mut comm = enclave_proc_spawn(&logger)
                .map_err(|err| {
                    err.add_subaction("Failed to spawn enclave process".to_string())
                        .set_action(RUN_ENCLAVE_STR.to_string())
                })
                .ok_or_exit_with_errno(None);

            let names = get_all_enclave_names()
                .map_err(|e| {
                    e.add_subaction("Failed to handle all enclave process replies".to_string())
                        .set_action("Get Enclaves Name".to_string())
                })
                .ok_or_exit_with_errno(None);
            run_args.enclave_name = Some(
                new_enclave_name(run_args.clone(), names)
                    .map_err(|err| {
                        err.add_subaction("Failed to assign a new enclave name".to_string())
                            .set_action(NEW_NAME_STR.to_string())
                    })
                    .ok_or_exit_with_errno(None),
            );

            enclave_proc_command_send_single(
                EnclaveProcessCommandType::Run,
                Some(&run_args),
                &mut comm,
            )
            .map_err(|e| {
                e.add_subaction("Failed to send single command".to_string())
                    .set_action(RUN_ENCLAVE_STR.to_string())
            })
            .ok_or_exit_with_errno(None);

            info!("Sent command: Run");
            replies.push(comm);
            let run_info = enclave_process_handle_all_replies::<EnclaveRunInfo>(
                &mut replies,
                0,
                false,
                vec![0],
            )
            .map_err(|e| {
                e.add_subaction("Failed to handle all enclave process replies".to_string())
                    .set_action(RUN_ENCLAVE_STR.to_string())
            })
            .ok_or_exit_with_errno(None);
            let enclave_cid = run_info
                .first()
                .map(|run_info| run_info.enclave_cid)
                .ok_or_else(|| {
                    new_nitro_cli_failure!(
                        "Enclave CID was not reported",
                        NitroCliErrorEnum::EnclaveConsoleConnectionFailure
                    )
                })
                .ok_or_exit_with_errno(None);
            if run_args.attach_console {
                console_enclaves(enclave_cid, None)
                    .map_err(|e| {
                        e.add_subaction("Failed to connect to enclave console".to_string())
                            .set_action(ENCLAVE_CONSOLE_STR.to_string())
                    })
                    .ok_or_exit_with_errno(None);
            }
        }
        Some(("terminate-enclave", args)) => {
            if args.is_present("all") {
                terminate_all_enclaves()
                    .map_err(|e| {
                        e.add_subaction("Failed to terminate all running enclaves".to_string())
                            .set_action(TERMINATE_ALL_ENCLAVES_STR.to_string())
                    })
                    .ok_or_exit_with_errno(None);
            } else {
                let terminate_args = TerminateEnclavesArgs::new_with(args)
                    .map_err(|err| {
                        err.add_subaction(
                            "Failed to construct TerminateEnclave arguments".to_string(),
                        )
                        .set_action(TERMINATE_ENCLAVE_STR.to_string())
                    })
                    .ok_or_exit_with_errno(None);
                let mut comm = enclave_proc_connect_to_single(&terminate_args.enclave_id)
                    .map_err(|e| {
                        e.add_subaction("Failed to connect to enclave process".to_string())
                            .set_action(TERMINATE_ENCLAVE_STR.to_string())
                    })
                    .ok_or_exit_with_errno(None);
                // TODO: Replicate output of old CLI on invalid enclave IDs.
                enclave_proc_command_send_single::<EmptyArgs>(
                    EnclaveProcessCommandType::Terminate,
                    None,
                    &mut comm,
                )
                .map_err(|e| {
                    e.add_subaction("Failed to send single command".to_string())
                        .set_action(TERMINATE_ENCLAVE_STR.to_string())
                })
                .ok_or_exit_with_errno(None);

                info!("Sent command: Terminate");
                replies.push(comm);
                enclave_process_handle_all_replies::<EnclaveTerminateInfo>(
                    &mut replies,
                    0,
                    false,
                    vec![0],
                )
                .map_err(|e| {
                    e.add_subaction("Failed to handle all enclave process replies".to_string())
                        .set_action(TERMINATE_ENCLAVE_STR.to_string())
                })
                .ok_or_exit_with_errno(None);
            }
        }
        Some(("describe-enclaves", args)) => {
            let describe_args = DescribeEnclavesArgs::new_with(args);
            let (comms, comm_errors) = enclave_proc_command_send_all::<DescribeEnclavesArgs>(
                EnclaveProcessCommandType::Describe,
                Some(&describe_args),
            )
            .map_err(|e| {
                e.add_subaction(
                    "Failed to send DescribeEnclave command to all enclave processes".to_string(),
                )
                .set_action(DESCRIBE_ENCLAVE_STR.to_string())
            })
            .ok_or_exit_with_errno(None);

            info!("Sent command: Describe");
            replies.extend(comms);
            enclave_process_handle_all_replies::<EnclaveDescribeInfo>(
                &mut replies,
                comm_errors,
                true,
                vec![0],
            )
            .map_err(|e| {
                e.add_subaction("Failed to handle all enclave process replies".to_string())
                    .set_action(DESCRIBE_ENCLAVE_STR.to_string())
            })
            .ok_or_exit_with_errno(None);
        }
        Some(("build-enclave", args)) => {
            let build_args = BuildEnclavesArgs::new_with(args)
                .map_err(|e| {
                    e.add_subaction("Failed to construct BuildEnclave arguments".to_string())
                        .set_action(BUILD_ENCLAVE_STR.to_string())
                })
                .ok_or_exit_with_errno(None);
            build_enclaves(build_args)
                .map_err(|e| {
                    e.add_subaction("Failed to build enclave".to_string())
                        .set_action(BUILD_ENCLAVE_STR.to_string())
                })
                .ok_or_exit_with_errno(None);
        }
        Some(("describe-eif", args)) => {
            let eif_path = args
                .value_of("eif-path")
                .map(|val| val.to_string())
                .unwrap();
            describe_eif(eif_path)
                .map_err(|e| {
                    e.add_subaction("Failed to describe EIF".to_string())
                        .set_action(DESCRIBE_EIF_STR.to_string())
                })
                .ok_or_exit_with_errno(None);
        }
        Some(("console", args)) => {
            let console_args = ConsoleArgs::new_with(args)
                .map_err(|e| {
                    e.add_subaction("Failed to construct Console arguments".to_string())
                        .set_action(ENCLAVE_CONSOLE_STR.to_string())
                })
                .ok_or_exit_with_errno(None);
            let enclave_cid = enclave_proc_get_cid(&console_args.enclave_id)
                .map_err(|e| {
                    e.add_subaction("Failed to retrieve enclave CID".to_string())
                        .set_action(ENCLAVE_CONSOLE_STR.to_string())
                })
                .ok_or_exit_with_errno(None);
            let enclave_flags = enclave_proc_get_flags(&console_args.enclave_id)
                .map_err(|e| {
                    e.add_subaction("Failed to retrieve enclave flags".to_string())
                        .set_action(ENCLAVE_CONSOLE_STR.to_string())
                })
                .ok_or_exit_with_errno(None);
            if enclave_flags & NE_ENCLAVE_DEBUG_MODE == 0 {
                let _result : NitroCliResult<()> = Err(new_nitro_cli_failure!(
                    "The enclave was not started with the debug flag set, include '--debug-mode' in the run-enclave command",
                    NitroCliErrorEnum::EnclaveConsoleConnectionFailure
                ))
                .map_err(|e| {
                    e.add_subaction("Failed to connect to enclave console".to_string())
                        .set_action(ENCLAVE_CONSOLE_STR.to_string())
                })
                .ok_or_exit_with_errno(None);
            }

            console_enclaves(enclave_cid, console_args.disconnect_timeout_sec)
                .map_err(|e| {
                    e.add_subaction("Failed to connect to enclave console".to_string())
                        .set_action(ENCLAVE_CONSOLE_STR.to_string())
                })
                .ok_or_exit_with_errno(None);
        }
        Some(("pcr", args)) => {
            let pcr_args = PcrArgs::new_with(args)
                .map_err(|e| {
                    e.add_subaction("Failed to construct PCR arguments".to_string())
                        .set_action(FILE_PCR_STR.to_string())
                })
                .ok_or_exit_with_errno(None);

            get_file_pcr(pcr_args.path, pcr_args.pcr_type)
                .map_err(|e| {
                    e.add_subaction("Failed to get the PCR hash of the file contents".to_string())
                        .set_action(FILE_PCR_STR.to_string())
                })
                .ok_or_exit_with_errno(None);
        }
        Some(("explain", args)) => {
            let explain_args = ExplainArgs::new_with(args)
                .map_err(|e| {
                    e.add_subaction("Failed to construct Explain arguments".to_string())
                        .set_action(EXPLAIN_ERR_STR.to_string())
                })
                .ok_or_exit_with_errno(None);
            explain_error(explain_args.error_code_str);
        }
        Some((&_, _)) | None => (),
    }
}


// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
#![deny(missing_docs)]
#![deny(warnings)]

use libc::{c_void, close};
use nix::poll::poll;
use nix::poll::{PollFd, PollFlags};
use nix::sys::socket::{connect, socket};
use nix::sys::socket::{AddressFamily, SockFlag, SockType, VsockAddr};
use nix::sys::time::{TimeVal, TimeValLike};
use nix::unistd::read;
use std::io::Write;
use std::mem::size_of;
use std::os::unix::io::IntoRawFd;
use std::os::unix::io::RawFd;
use std::thread::sleep;
use std::time::{Duration, SystemTime};
use vmm_sys_util::epoll::{ControlOperation, Epoll, EpollEvent, EventSet};
use vmm_sys_util::timerfd::TimerFd;

use crate::common::{NitroCliErrorEnum, NitroCliFailure, NitroCliResult};
use crate::new_nitro_cli_failure;

/// The size of the buffers used for reading console data.
const BUFFER_SIZE: usize = 1024;

/// The console connection time-out, in milliseconds.
const CONSOLE_CONNECT_TIMEOUT: i64 = 20000;

/// The `poll` time-out, in milliseconds.
const POLL_TIMEOUT: i32 = 10000;

/// The socket connection time-out flag.
const SO_VM_SOCKETS_CONNECT_TIMEOUT: i32 = 6;

/// The amount of time to wait between consecutive console reads, in milliseconds.
const TIMEOUT: u64 = 100;

/// Defines the types of PCRs that can be measured by `pcr` command
pub enum PcrType {
    /// Used for files containing the bytes for hashing
    DefaultType,
    /// Used for `.pem` files that we want to hash. Additional serializing is needed
    SigningCertificate,
}

/// The structure representing the console of an enclave.
pub struct Console {
    /// The file descriptor used for connecting to the enclave's console.
    fd: RawFd,
}

impl Drop for Console {
    fn drop(&mut self) {
        unsafe { close(self.fd) };
    }
}

impl Console {
    /// Create a new blocking `Console` connection from a given enclave CID and a vsock port.
    pub fn new(cid: u32, port: u32) -> NitroCliResult<Self> {
        let socket_fd = socket(
            AddressFamily::Vsock,
            SockType::Stream,
            SockFlag::empty(),
            None,
        )
        .map_err(|err| {
            new_nitro_cli_failure!(
                &format!("Failed to create blocking console socket: {:?}", err),
                NitroCliErrorEnum::SocketError
            )
        })?;

        let sockaddr = VsockAddr::new(cid, port);

        vsock_set_connect_timeout(socket_fd, CONSOLE_CONNECT_TIMEOUT).map_err(|err| {
            err.add_subaction("Failed to set console connect timeout".to_string())
        })?;

        connect(socket_fd, &sockaddr).map_err(|err| {
            new_nitro_cli_failure!(
                &format!("Failed to connect to the console: {:?}", err),
                NitroCliErrorEnum::EnclaveConsoleConnectionFailure
            )
        })?;

        Ok(Console { fd: socket_fd })
    }

    /// Create a new non-blocking `Console` connection from a given enclave CID and a vsock port.
    pub fn new_nonblocking(cid: u32, port: u32) -> NitroCliResult<Self> {
        // create new non blocking socket
        let socket_fd = socket(
            AddressFamily::Vsock,
            SockType::Stream,
            SockFlag::SOCK_NONBLOCK,
            None,
        )
        .map_err(|err| {
            new_nitro_cli_failure!(
                &format!("Failed to create nonblocking console socket: {:?}", err),
                NitroCliErrorEnum::SocketError
            )
        })?;

        vsock_set_connect_timeout(socket_fd, CONSOLE_CONNECT_TIMEOUT).map_err(|err| {
            err.add_subaction("Failed to set console connect timeout".to_string())
        })?;

        let sockaddr = VsockAddr::new(cid, port);
        let result = connect(socket_fd, &sockaddr);

        match result {
            Ok(_) => println!("Connected to the console"),
            Err(error) => match error {
                // If the connection is not ready, wait until socket_fd is ready for writing.
                nix::errno::Errno::EINPROGRESS => {
                    let poll_fd = PollFd::new(socket_fd, PollFlags::POLLOUT);
                    let mut poll_fds = [poll_fd];
                    match poll(&mut poll_fds, POLL_TIMEOUT) {
                        Ok(1) => println!("Connected to the console"),
                        _ => {
                            return Err(new_nitro_cli_failure!(
                                "Failed to connect to the console",
                                NitroCliErrorEnum::SocketError
                            ))
                        }
                    }
                }
                _ => {
                    return Err(new_nitro_cli_failure!(
                        "Failed to connect to the console",
                        NitroCliErrorEnum::SocketError
                    ))
                }
            },
        };

        Ok(Console { fd: socket_fd })
    }

    /// Read a chunk of raw data from the console and output it.
    pub fn read_to(
        &self,
        output: &mut dyn Write,
        disconnect_timeout_sec: Option<u64>,
    ) -> NitroCliResult<()> {
        // Initialize variables
        let epoll = Epoll::new().map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to create epoll: {:?}", e),
                NitroCliErrorEnum::EpollError
            )
        })?;

        // Add console fd to epoll
        epoll
            .ctl(
                ControlOperation::Add,
                self.fd,
                EpollEvent::new(EventSet::IN, self.fd as u64),
            )
            .map_err(|e| {
                new_nitro_cli_failure!(
                    &format!("Failed to add fd to epoll: {:?}", e),
                    NitroCliErrorEnum::EpollError
                )
            })?;

        // If the function call provides a disconnect timeout, create a timerfd,
        // arm it and then add it to epoll
        if let Some(disconnect_timeout) = disconnect_timeout_sec {
            // Create timerfd
            let mut timerfd = TimerFd::new().map_err(|e| {
                new_nitro_cli_failure!(
                    &format!("Failed to initialize timerfd: {:?}", e),
                    NitroCliErrorEnum::EpollError
                )
            })?;

            // Arm timerfd with disconnect_timeout seconds
            timerfd
                .reset(Duration::from_secs(disconnect_timeout), None)
                .map_err(|e| {
                    new_nitro_cli_failure!(
                        &format!("Failed to arm timerfd: {:?}", e),
                        NitroCliErrorEnum::EpollError
                    )
                })?;

            // Add timerfd fd to epoll
            let timerfd_fd = timerfd.into_raw_fd();
            epoll
                .ctl(
                    ControlOperation::Add,
                    timerfd_fd,
                    EpollEvent::new(EventSet::IN, timerfd_fd as u64),
                )
                .map_err(|e| {
                    new_nitro_cli_failure!(
                        &format!("Failed to add fd to epoll: {:?}", e),
                        NitroCliErrorEnum::EpollError
                    )
                })?;
        }

        // Allow only one epoll event to happen at a given time
        let mut events = [EpollEvent::default(); 1];

        loop {
            // Wait for kernel notification that one of the fds is available
            let num_events = epoll.wait(-1, &mut events).map_err(|e| {
                new_nitro_cli_failure!(
                    &format!("Failed to wait epoll: {:?}", e),
                    NitroCliErrorEnum::EpollError
                )
            })?;

            // Check if any event triggered, because an interrupt could unblock the wait
            // without any of the requested events to occur
            if num_events == 1 {
                match events[0].fd() {
                    // Check if console fd triggered
                    fd if fd == self.fd => {
                        let mut buffer = [0u8; BUFFER_SIZE];
                        let size = read(self.fd, &mut buffer).map_err(|e| {
                            new_nitro_cli_failure!(
                                &format!("Failed to read data from the console: {:?}", e),
                                NitroCliErrorEnum::EnclaveConsoleReadError
                            )
                        })?;

                        if size == 0 {
                            break;
                        }

                        if size > 0 {
                            output.write(&buffer[..size]).map_err(|e| {
                                new_nitro_cli_failure!(
                                    &format!(
                                        "Failed to write data from the \
                                        console to the given stream: {:?}",
                                        e
                                    ),
                                    NitroCliErrorEnum::EnclaveConsoleWriteOutputError
                                )
                            })?;
                        }
                    }
                    // Check if timerfd triggered
                    _ => break,
                }
            }
        }

        Ok(())
    }

    /// Read a chunk of raw data to a buffer.
    pub fn read_to_buffer(&self, buf: &mut Vec<u8>, duration: Duration) -> NitroCliResult<()> {
        let sys_time = SystemTime::now();

        loop {
            let mut buffer = [0u8; BUFFER_SIZE];
            let result = read(self.fd, &mut buffer);

            if let Ok(size) = result {
                if size > 0 {
                    let mut buf_vec = buffer.to_vec();
                    buf_vec.truncate(size);
                    (*buf).append(&mut buf_vec);
                }
            }

            sleep(Duration::from_millis(TIMEOUT));

            let time_elapsed = sys_time.elapsed().map_err(|err| {
                new_nitro_cli_failure!(
                    &format!("System time moved backwards: {:?}", err),
                    NitroCliErrorEnum::ClockSkewError
                )
            })?;

            if time_elapsed >= duration {
                break;
            }
        }

        Ok(())
    }
}

/// Set a timeout on a vsock connection.
fn vsock_set_connect_timeout(fd: RawFd, millis: i64) -> NitroCliResult<()> {
    let timeval = TimeVal::milliseconds(millis);
    let ret = unsafe {
        libc::setsockopt(
            fd,
            libc::AF_VSOCK,
            SO_VM_SOCKETS_CONNECT_TIMEOUT,
            &timeval as *const _ as *const c_void,
            size_of::<TimeVal>() as u32,
        )
    };

    match ret {
        0 => Ok(()),
        _ => Err(new_nitro_cli_failure!(
            &format!(
                "Failed to configure SO_VM_SOCKETS_CONNECT_TIMEOUT: {:?}",
                ret
            ),
            NitroCliErrorEnum::SocketConnectTimeoutError
        )),
    }
}

/// Computes the ceil of `lhs / rhs`. Used for reporting the lower
/// limit of enclave memory based on the EIF file size.
pub fn ceil_div(lhs: u64, rhs: u64) -> u64 {
    if rhs == 0 {
        return std::u64::MAX;
    }

    lhs / rhs
        + match lhs % rhs {
            0 => 0,
            _ => 1,
        }
}