#![no_std]
#![no_main]
#![feature(abi_efiapi)]
#![feature(asm)]
#![feature(maybe_uninit_extra)]
#![deny(unsafe_op_in_unsafe_fn)]
#![feature(default_alloc_error_handler)]

#[macro_use]
extern crate uefi;
#[macro_use]
extern crate alloc;
#[macro_use]
extern crate log;

use alloc::boxed::Box;
use uefi::{Handle, Status, prelude::{entry, Boot, SystemTable}, proto::console::gop::PixelFormat, Completion, ResultExt};
use log::*;
use uefi::table::boot::{AllocateType, MemoryDescriptor, MemoryType};
use alloc::vec::*;
use core::mem::transmute;
use core::panic::PanicInfo;
use aarch64::asm::get_current_el;
use aarch64::cpu::FieldValue;
use aarch64::PhysAddr;
use aarch64::regs::RegisterReadWrite;
use acpi::{AcpiTables, PlatformInfo};
use uefi::prelude::*;
use uefi::proto::{
    console::gop::GraphicsOutput,
    media::file::{
        FileType::Regular,
        {File, FileAttribute, FileInfo, FileMode, RegularFile}
    }
};
use xmas_elf::ElfFile;
use xmas_elf::program::{ProgramHeader, ProgramHeader64, Type};
use uefi::proto::media::file::Directory;
use bootloader::arch::aarch64::uefi::*;
use libsm::sm2::signature::*;
use libsm::sm3::hash::Sm3Hash;
use uefi::table::cfg::ACPI2_GUID;

/*
#[link(name="optee", kind="static")]
extern "efiapi" {
    pub fn IsOpteePresent() -> bool;
}
 */

#[repr(C)]
#[derive(Debug)]
struct kernel_header {
    pub_key: [u8; 65],
    sign_r: [u8; 32],
    sign_s: [u8; 32]
}

const KERNEL_LOCATION: &'static str = "kernel";
const PUBLIC_KEY_HASH: &'static str = "pk_hash";
const PAGE_SIZE: usize = 0x1000;
const PAGE_SIZE_BITS: usize = 12;

#[entry]
fn efi_main(image: Handle, st: SystemTable<Boot>) -> Status {
    // Initialize utilities (logging, memory allocation...)
    uefi_services::init(&st).expect_success("Failed to initialize utilities");
    let boot_services = st.boot_services();
    boot_services.set_watchdog_timer(0, 0x10000, None)
        .expect_success("Could not disable UEFI watchdog");
    st.stdout().clear().expect_success("unable to clear screen");

    let fs = st.
        boot_services().
        get_image_file_system(image.clone()).
        expect_success("cannot get image file system");
    let fs = unsafe { &mut *fs.get() };
    let mut root_dir = fs.open_volume().expect_success("cannot get root dir");

    info!("efi_main address: 0x{:x}", efi_main as usize);
    info!("Current EL: {}", get_current_el());

    let rsdp_addr = {
        use uefi::table::cfg;
        let mut config_entries = st.config_table().iter();
        for conf_entry in config_entries.as_ref() {
            info!("config entry: {:?}", conf_entry);
        }
        // look for an ACPI2 RSDP first
        let acpi2_rsdp = config_entries.find(|entry| matches!(entry.guid, cfg::ACPI2_GUID));
        // if no ACPI2 RSDP is found, look for a ACPI1 RSDP
        let rsdp = acpi2_rsdp
            .or_else(|| config_entries.find(|entry| matches!(entry.guid, cfg::ACPI_GUID)));
        rsdp.map(|entry| PhysAddr::new(entry.address as u64))
    }.expect("rsdp not found");
    info!("rsdp_addr: {:?}", rsdp_addr);

    let mut kernel_img = root_dir.open(KERNEL_LOCATION, FileMode::Read, FileAttribute::READ_ONLY).expect_success("cannot open kernel image");
    let mut kernel_img = unsafe { RegularFile::new(kernel_img) };
    let kernel_info: Box<FileInfo> = kernel_img.get_boxed_info().unwrap_success();
    let kernel_size = kernel_info.file_size();

    // let mut kernel_buf = vec![0 as u8; kernel_size as usize];
    // kernel_img.read(kernel_buf.as_mut_slice());
    // let start_addr = 0x40000000;
    // let pages = kernel_size as usize / PAGE_SIZE + 1;
    // boot_services.allocate_pages(AllocateType::Address(start_addr), MemoryType::LOADER_DATA, pages);
    // let dst = unsafe { core::slice::from_raw_parts_mut(start_addr as *mut u8, kernel_size as usize) };
    // let src = kernel_buf.as_slice();
    // dst.copy_from_slice(src);

    info!("running kernel integrity check...");
    let kernel_elf = verify_kernel(&mut root_dir, &mut kernel_img, kernel_size as usize);
    info!("loading kernel to memory...");
    let kernel_entry = load_kernel(boot_services, kernel_elf);
    info!("kernel entry: 0x{:x}", kernel_entry);
    let kernel_entry: extern "C" fn(usize, usize, usize) = unsafe { transmute(kernel_entry) };

    let mmap_storage = {
        let max_mmap_size =
            boot_services.memory_map_size() + 8 * core::mem::size_of::<MemoryDescriptor>();
        let ptr = boot_services
            .allocate_pool(MemoryType::LOADER_DATA, max_mmap_size)?
            .log();
        unsafe { core::slice::from_raw_parts_mut(ptr, max_mmap_size) }
    };

    info!("current pc: 0x{:x}", unsafe { aarch64::asm::get_pc() });
    info!("exit boot services");
    let (system_table, memory_map) = st
        .exit_boot_services(image, mmap_storage)
        .expect_success("Failed to exit boot services");

    kernel_entry(0, 0, 0);

    Status::SUCCESS
}

fn load_kernel(boot_services: &BootServices, kernel_elf: Vec<u8>) -> u64 {
    let kernel_elf = xmas_elf::ElfFile::new(kernel_elf.as_slice()).unwrap();
    let elf_header = kernel_elf.header;
    info!("ELF header info: {:?}", elf_header);

    for ph in kernel_elf.program_iter() {
        if ph.get_type().unwrap() == Type::Load {
            let start_va = ph.virtual_addr() as usize;
            let end_va = (ph.virtual_addr() + ph.mem_size()) as usize;
            let pages = (end_va >> PAGE_SIZE_BITS) - (start_va >> PAGE_SIZE_BITS) + 1;
            let ph_flags = ph.flags();
            info!("load header to address: 0x{:x}", start_va);

            boot_services.allocate_pages(AllocateType::Address(start_va), if ph_flags.is_write() { MemoryType::LOADER_DATA } else { MemoryType::LOADER_CODE }, pages);

            let dst = unsafe { core::slice::from_raw_parts_mut(start_va as *mut u8, ph.mem_size() as usize) };
            let src = &kernel_elf.input[ph.offset() as usize..(ph.offset() + ph.mem_size()) as usize];
            dst.copy_from_slice(src);

            info!("load header to address: 0x{:x} successfully", start_va);
        }
    }

    kernel_elf.header.pt2.entry_point()
}

fn verify_kernel(root_dir: &mut Directory, kernel_img: &mut RegularFile, kernel_size: usize) -> Vec<u8> {
    // load packaged kernel from disk and reset cursor to use behind
    let mut kernel_data = vec![0 as u8; kernel_size];
    let kernel_data = kernel_data.as_mut_slice();
    kernel_img.read(kernel_data);

    // split header and true kernel ELF
    let header = unsafe { (kernel_data.as_ptr() as *const kernel_header).as_ref() }.unwrap();
    let kernel_data = &kernel_data[core::mem::size_of::<kernel_header>()..];
    let elf_file = xmas_elf::ElfFile::new(kernel_data).unwrap();
    let elf_header = elf_file.header;
    let magic = elf_header.pt1.magic;
    assert_eq!(magic, [0x7f, 0x45, 0x4c, 0x46], "invalid elf!");

    // load public key from image and calculate the hash value
    let ctx = SigCtx::new();
    let pub_key = ctx.load_pubkey(&header.pub_key).expect("failed to load public key");
    let pk_hash = get_hash(&header.pub_key);

    // load public key hash from disk
    let mut pk_hash_file = root_dir.open(PUBLIC_KEY_HASH, FileMode::Read, FileAttribute::READ_ONLY).expect_success("unable to open pub_key_hash file");
    let pub_key_hash_info: Box<FileInfo> = pk_hash_file.get_boxed_info().unwrap().unwrap();
    let mut pk_hash_file = unsafe { RegularFile::new(pk_hash_file) };
    let mut pk_hash_buf = vec![0 as u8; pub_key_hash_info.file_size() as usize];
    let pk_hash_buf = pk_hash_buf.as_mut_slice();
    pk_hash_file.read(pk_hash_buf).expect_success("fail to read pub_key_hash");

    // check if both value of calc and value of disk are equal
    if !pk_hash_buf.eq(&pk_hash) {
        panic!("public key verification failed");
    } else {
        info!("public key verification passed");
    }

    // check the kernel integrity by hash and public key
    let kernel_hash = get_hash(kernel_data);
    let sign = Signature::new(&header.sign_r, &header.sign_s);

    if !ctx.verify(&kernel_hash, &pub_key, &sign) {
        panic!("kernel verification failed");
    } else {
        info!("kernel verification passed");
    }

    info!("kernel integrity checked pass!");

    kernel_data.to_vec()
}

fn get_hash(data: &[u8]) -> [u8; 32] {
    let mut sm3hasher = Sm3Hash::new(data);
    sm3hasher.get_hash()
}

#[panic_handler]
fn panic(_info: &PanicInfo) -> ! {
    info!("panic: {:?}", _info);
    loop {
        aarch64::asm::wfe();
    }
}
