// SPDX-License-Identifier: MPL-2.0
#include <cstdint>
#include <fstream>
#include <iostream>
#include <filesystem>
#include <string>
#include <vector>
#include <cstdlib>
#include <random>
#include <unistd.h>
#include <limits.h>
#include "elfio.hpp"
#include "elfio/elf_types.hpp"
using namespace std;
namespace fs=filesystem;
const string entry_flags="-ffreestanding -nostdlib -nostartfiles -Wall -Wextra -fno-stack-protector -fno-builtin -fno-pic -m64 -mcmodel=large";
struct vyx_addresses {
  uint64_t text_offset;
  uint64_t stack_base;
};
uint64_t rand_in_range(uint64_t min, uint64_t max) {
  static random_device rd;
  static mt19937_64 gen(rd());
  std::uniform_int_distribution<uint64_t> dis(min, max);
  return dis(gen);
}
uint64_t random_canonical(uint64_t start,uint64_t end) {
  uint64_t addr=rand_in_range(start/4096,end/4096)*4096;
  return addr;
}
vyx_addresses generate_vyx_addresses() {
  vyx_addresses addr;
  addr.text_offset=random_canonical(0xFFFF800000000000ULL,0xFFFF900000000000ULL-1);
  uint64_t stack_window_start=0xFFFFF00000000000ULL;
  uint64_t stack_window_end=0xFFFFFF8000000000ULL-0x100000ULL;
  addr.stack_base=random_canonical(stack_window_start,stack_window_end);
  return addr;
}
size_t pad_size(size_t value,size_t multiple) {
  if (multiple==0) return value;
  size_t remainder=value%multiple;
  if (remainder==0) return value;
  return value+multiple-remainder;
}
int main(int argc,char **argv) {
    cout<<"[Vyld] Vystem Executable Linker v0.1"<<endl;
    if (argc<3) {
        cout<<"[Vyld] Error: Wrong number of arguments."<<endl;
        return -1;
    }
    char buf[PATH_MAX];
    ssize_t len=readlink("/proc/self/exe",buf,sizeof(buf)-1);
    fs::path exe_path;
    if (len!=-1) {
        buf[len]='\0';
        exe_path=fs::path(buf);
    } else {
        cout<<"[Vyld] Error: couldn't obtain vyld executable path."<<endl;
        return -1;
    }
    exe_path=exe_path.parent_path();
    fs::path res_file=exe_path/fs::path("res.txt");
    fs::path combined_o=exe_path/fs::path("combined.o");
    fs::path compiled_elf=exe_path/fs::path("compiled.elf");
    fs::path script_ld=exe_path/fs::path("script.ld");
    if (!fs::exists(exe_path/fs::path("_vyx_start.c"))) {
        cout<<"[Vyld] Error: couldn't find entry point."<<endl;
        return -1;
    }
    const char* flags_c=getenv("VYLD_COMPILATION_FLAGS");
    string flags;
    if (flags_c) {
        flags=string(flags_c);
        if (flags!="") {
            cout<<"[Vyld] Using following flags: "<<flags<<endl;
        } else {
            cout<<"[Vyld] Warning: VYLD_COMPILATION_FLAGS is defined but empty."<<endl;
        }
    } else {
        flags="";
        cout<<"[Vyld] Warning: VYLD_COMPILATION_FLAGS isn't defined."<<endl;
    }
    cout<<"[Vyld] Compiling _vyx_start.c "<<endl;
    fs::remove(exe_path/"_vyx_start.o");
    int ret=system(string("gcc -c "+string(exe_path/"_vyx_start.c")+" -o "+string(exe_path/"_vyx_start.o")+" "+entry_flags+" > "+string(res_file)).c_str());
    if (ret!=0) {
        cout<<"[Vyld] Error: gcc didn't returned success. Error code: "<<ret<<endl;
        int ret=system(string("cat "+string(res_file)).c_str());
        return -1;
    } else {
        fs::remove(string(res_file));
    }
    vector<string> c_files(argv+1,argv+argc-1);
    vector<string> o_files;
    cout<<"[Vyld] Finding provided files:"<<endl;
    for (int i=0;i<c_files.size();i++) {
        if (!fs::exists(c_files[i])) {
            cout<<"[Vyld]   Error: can't found file \""<<c_files[i]<<"\"."<<endl;
            return -1;
        }
        string filename=fs::path(c_files[i]).stem().string();
        o_files.push_back(exe_path/fs::path(filename+".o"));
        cout<<"[Vyld]   Found \""<<fs::absolute(fs::path(c_files[i])).string()<<"\" that will be compile into \""<<o_files[i]<<"\" ("<<i+1<<"/"<<c_files.size()<<")."<<endl;
    }
    cout<<"[Vyld] Compiling provided files:"<<endl;
    for (int i=0;i<c_files.size();i++) {
        int ret=system(string("gcc -c "+c_files[i]+" -fno-pic -m64 -mcmodel=large "+flags+" -o "+o_files[i]+" > "+string(res_file)).c_str());
        if (ret!=0) {
            cout<<"[Vyld]   Error: gcc didn't returned success. File \""<<c_files[i]<<"\" returned error code: "<<ret<<endl;
            int ret=system(string("cat "+string(res_file)).c_str());
            return -1;
        } else {
            fs::remove(string(res_file));
        }
        cout<<"[Vyld]   Successfully compiled \""<<fs::absolute(fs::path(c_files[i])).string()<<"\" into \""<<o_files[i]<<"\" ("<<i+1<<"/"<<c_files.size()<<")."<<endl;
    }
    cout<<"[Vyld] Compiling all objects files into one...";
    string command_line_gcc="gcc -r "+string(exe_path/"_vyx_start.o")+" ";
    for (auto o:o_files) {
        command_line_gcc+=o+" ";
    }
    command_line_gcc+="-o "+string(combined_o)+" > "+string(res_file);
    ret=system(command_line_gcc.c_str());
    if (ret!=0) {
        cout<<endl<<"[Vyld] Error: gcc didn't returned success. Error code: "<<ret<<endl;
        int ret=system(string("cat "+string(res_file)).c_str());
        return -1;
    }
    cout<<" done."<<endl;
    cout<<"[Vyld] Opening combined object files...";
    ELFIO::elfio reader;
    if (!reader.load(string(combined_o))) {
        cout<<endl<<"[Vyld] Error: couldn't open combined.o."<<endl;
        return -1;
    }
    if (reader.get_type()!=ELFIO::ET_REL) {
        cout<<endl<<"[Vyld] Error: combined.o isn't a relocatable object file."<<endl;
        return -1;
    }
    if (reader.get_class()!=ELFIO::ELFCLASS64) {
        cout<<endl<<"[Vyld] Error: combined.o isn't a 64-bit ELF file."<<endl;
        return -1;
    }
    if (reader.get_machine()!=ELFIO::EM_X86_64) {
        cout<<endl<<"[Vyld] Error: combined.o isn't a x86_64 file."<<endl;
        return -1;
    }
    cout<<" done."<<endl;
    cout<<"[Vyld] Extracting important sections."<<endl;
    const ELFIO::section* text_sec=reader.sections[".text"];
    const ELFIO::section* data_sec=reader.sections[".data"];
    const ELFIO::section* rodata_sec=reader.sections[".rodata"];
    const ELFIO::section* bss_sec=reader.sections[".bss"];
    if (!text_sec) {
        cout<<endl<<"[Vyld] Error: no .text found."<<endl;
        return -1;
    }
    if (!data_sec) {
        cout<<"[Vyld] Warning: no .data found."<<endl;
    }
    if (!rodata_sec) {
        cout<<"[Vyld] Warning: no .rodata found."<<endl;
    }
    if (!bss_sec) {
        cout<<"[Vyld] Warning: no .bss found."<<endl;
    }
    if (text_sec->get_size()==0) {
        cout<<endl<<"[Vyld] Error: .text is empty."<<endl;
        return -1;
    }
    cout<<"[Vyld] Generating linker script:"<<endl;
    size_t text_size=text_sec->get_size();
    size_t text_size_padded=pad_size(text_size,4096);
    cout<<"[Vyld] text_size=0x"<<hex<<text_size<<", text_size_padded=0x"<<hex<<text_size_padded<<endl;
    size_t data_size=0;
    size_t data_size_padded=0;
    size_t rodata_size=0;
    size_t rodata_size_padded=0;
    size_t bss_size=0;
    size_t bss_size_padded=0;
    if (data_sec) {
        data_size=data_sec->get_size();
        data_size_padded=pad_size(data_size,4096);
        cout<<"[Vyld] data_size=0x"<<hex<<data_size<<", data_size_padded=0x"<<hex<<data_size_padded<<endl;
    }
    if (rodata_sec) {
        rodata_size=rodata_sec->get_size();
        for (const auto &s:reader.sections) {
          if (s->get_name().size()>7 && s->get_name().substr(0,7)==".rodata" && s->get_name()!=".rodata") {
            rodata_size+=s->get_size();
          }
        }
        rodata_size_padded=pad_size(rodata_size,4096);
        cout<<"[Vyld] rodata_size=0x"<<hex<<rodata_size<<", rodata_size_padded=0x"<<hex<<rodata_size_padded<<endl;
    }
    if (bss_sec) {
        bss_size=bss_sec->get_size();
        bss_size_padded=pad_size(bss_size,4096);
        cout<<"[Vyld] bss_size=0x"<<hex<<bss_size<<", bss_size_padded=0x"<<hex<<bss_size_padded<<endl;
    }
    vyx_addresses va=generate_vyx_addresses();
    size_t text_offset=va.text_offset;
    cout<<"[Vyld] Loading address (va.text_offset): 0x"<<hex<<text_offset<<endl;
    cout<<"[Vyld] Stack address base (va.stack_base): 0x"<<hex<<va.stack_base<<endl;
    size_t data_offset=text_offset+text_size_padded;
    size_t rodata_offset=data_offset+data_size_padded;
    size_t bss_offset=rodata_offset+rodata_size_padded;
    ofstream ld(script_ld);
    ld<<"ENTRY(_vyx_start)\n";
    ld<<"MEMORY\n{\n";
    ld<<"  TEXT (rx) : ORIGIN = 0x"<<hex<<text_offset<<", LENGTH = "<<dec<<text_size_padded<<"\n";
    ld<<"  DATA (rw) : ORIGIN = 0x"<<hex<<data_offset<<", LENGTH = "<<dec<<data_size_padded<<"\n";
    ld<<"  RODATA (r) : ORIGIN = 0x"<<hex<<rodata_offset<<", LENGTH = "<<dec<<rodata_size_padded<<"\n";
    ld<<"  BSS (rw) : ORIGIN = 0x"<<hex<<bss_offset<<", LENGTH = "<<dec<<bss_size_padded<<"\n";
    ld<<"}\nSECTIONS\n{\n";
    ld<<"  .text :\n";
    ld<<"  {\n";
    ld<<"    . = ORIGIN(TEXT);\n";
    ld<<"    *(.text)\n";
    ld<<"    *(.text.*)\n";
    ld<<"    . = ALIGN(0x1000);\n";
    ld<<"  } > TEXT\n\n";
    ld<<"  .data :\n";
    ld<<"  {\n";
    ld<<"    . = ORIGIN(DATA);\n";
    ld<<"    *(.data)\n";
    ld<<"    *(.data.*)\n";
    ld<<"    . = ALIGN(0x1000);\n";
    ld<<"  } > DATA\n\n";
    ld<<"  .rodata :\n";
    ld<<"  {\n";
    ld<<"    . = ORIGIN(RODATA);\n";
    ld<<"    *(.rodata)\n";
    ld<<"    *(.rodata.*)\n";
    ld<<"    . = ALIGN(0x1000);\n";
    ld<<"  } > RODATA\n\n";
    ld<<"  .bss :\n";
    ld<<"  {\n";
    ld<<"    . = ORIGIN(BSS);\n";
    ld<<"    *(.bss)\n";
    ld<<"    *(.bss.*)\n";
    ld<<"    . = ALIGN(0x1000);\n";
    ld<<"  } > BSS\n\n";
    ld<<"  /DISCARD/ :\n";
    ld<<"  {\n";
    ld<<"    *(.note.*)\n";
    ld<<"    *(.comment)\n";
    ld<<"    *(.eh_frame)\n";
    ld<<"    *(.eh_frame_hdr)\n";
    ld<<"  }\n}";
    ld.close();
    cout<<"[Vyld] Linking relocatable object."<<endl;
    ret=system(string("ld -T "+string(script_ld)+" -o "+string(compiled_elf)+" "+string(combined_o)).c_str());
    if (ret!=0) {
        cout<<endl<<"[Vyld] Error: ld didn't returned success. Error code: "<<ret<<endl;
        int ret=system(string("cat res.txt").c_str());
        return -1;
    }
    cout<<"[Vyld] Opening compiled.elf."<<endl;
    if (!reader.load(string(compiled_elf))) {
        cout<<"[Vyld] Error: couldn't open compiled.elf."<<endl;
        return -1;
    }
    if (reader.get_type()!=ELFIO::ET_EXEC) {
        cout<<"[Vyld] Error: compiled.elf isn't a relocatable object file."<<endl;
        return -1;
    }
    if (reader.get_class()!=ELFIO::ELFCLASS64) {
        cout<<"[Vyld] Error: compiled.elf isn't a 64-bit ELF file."<<endl;
        return -1;
    }
    if (reader.get_machine()!=ELFIO::EM_X86_64) {
        cout<<"[Vyld] Error: compiled.elf isn't a x86_64 file."<<endl;
        return -1;
    }
    cout<<"[Vyld] Extracting linked sections from compiled.elf."<<endl;
    const ELFIO::section* elf_text_sec=reader.sections[".text"];
    const ELFIO::section* elf_data_sec=reader.sections[".data"];
    const ELFIO::section* elf_rodata_sec=reader.sections[".rodata"];
    const ELFIO::section* elf_bss_sec=reader.sections[".bss"];
    if (!elf_text_sec) {
        cout<<"[Vyld] Error: no .text found."<<endl;
        return -1;
    }
    if (!elf_data_sec) {
        cout<<"[Vyld] Warning: no .data found."<<endl;
    }
    if (!elf_rodata_sec) {
        cout<<"[Vyld] Warning: no .rodata found."<<endl;
    }
    if (!elf_bss_sec) {
        cout<<"[Vyld] Warning: no .bss found."<<endl;
    }
    if (elf_text_sec->get_size()==0) {
        cout<<"[Vyld] Error: .text is empty."<<endl;
        return -1;
    }
    vector<unsigned char> elf_text_bin(text_size_padded,0);
    vector<unsigned char> elf_data_bin(data_size_padded,0);
    vector<unsigned char> elf_rodata_bin(rodata_size_padded,0);
    memcpy(elf_text_bin.data(),elf_text_sec->get_data(),elf_text_sec->get_size());
    if (elf_data_sec->get_size()!=0) {
      memcpy(elf_data_bin.data(),elf_data_sec->get_data(),elf_data_sec->get_size());
    }
    if (elf_rodata_sec->get_size()!=0) {
      memcpy(elf_rodata_bin.data(),elf_rodata_sec->get_data(),elf_rodata_sec->get_size());
    }
    uint8_t vyx_sig[3]={'V','y','X'};
    uint16_t vyx_ver=0x0001;
    string outfile=string(argv[argc-1]);
    ofstream out(outfile,ios::binary);
    out.write(reinterpret_cast<char*>(vyx_sig),3);
    out.write(reinterpret_cast<char*>(&vyx_ver),2);
    out.write(reinterpret_cast<char*>(&va.text_offset),8);
    out.write(reinterpret_cast<char*>(&va.stack_base),8);
    out.write(reinterpret_cast<char*>(&text_size_padded),8);
    out.write(reinterpret_cast<char*>(&data_size_padded),8);
    out.write(reinterpret_cast<char*>(&rodata_size_padded),8);
    uint64_t bss_size_elf=elf_bss_sec->get_size();
    out.write(reinterpret_cast<char*>(&bss_size_elf),8);
    out.write(reinterpret_cast<char*>(elf_text_bin.data()),elf_text_bin.size());
    if (data_size_padded!=0) out.write(reinterpret_cast<char*>(elf_data_bin.data()),elf_data_bin.size());
    if (rodata_size_padded!=0) out.write(reinterpret_cast<char*>(elf_rodata_bin.data()),elf_rodata_bin.size());
    out.close();
    cout<<"[Vyld] Succesfully compiled "<<outfile<<endl;
    return 0;
}