local IMAGE_SCN_CNT_CODE=0x20
local IMAGE_SCN_MEM_EXECUTE=0x20000000

local gddnAddress=nil
local gddnName=nil
function getDriverName(address) --not multithreaded so watch it
  if gddnAddress==nil then
    local s1=[[
alloc(gddn,512)
alloc(name,256)

gddn:
  sub rsp,28
  mov rdx,name
  mov r8,#256
  call GetDevicedriverBaseNameA
  add rsp,28
  ret
  ]]

    local r,info=autoAssemble(s1,true)
    if r then
      gddnName=info.allocs.name.address
      gddnAddress=info.allocs.gddn.address
    else
      error('getDriverName autoAssemble failed:'..info)
    end
  end

  local Result=executeCodeLocal(gddnAddress,address)

  if Result>0 then
    local Name=readStringLocal(gddnName,Result)
    return Name
  end
end

local gddpAddress=nil
local gddpFilename=nil
function getDriverPath(address)
  if gddpAddress==nil then
    local s1=[[
alloc(gddp,512)
alloc(filename,256)

gddp:
  sub rsp,28
  mov rdx,filename
  mov r8,#256
  call GetDeviceDriverFileNameA
  add rsp,28
  ret
  ]]

    local r
    local info
    r,info=autoAssemble(s1,true)
    if r then
      gddpFilename=info.allocs.filename.address
      gddpAddress=info.allocs.gddp.address
    else
      error('getDriverPath autoAssemble failed:'..info)
    end
  end

  local Result=executeCodeLocal(gddpAddress,address)

  if Result>0 then
    local Path=readStringLocal(gddpFilename,Result)
    return Path
  end
end

function getDriverList()
  local s1=[[
[enable]
//loadlibrary(psapi.dll)
alloc(enumddp1,512)
alloc(bytesneeded,8)
alloc(done,1)

enumddp1:
sub rsp,28
mov rcx,0
mov rdx,0
mov r8,bytesneeded
call k32EnumDeviceDrivers
add rsp,28

mov byte [done],1
ret

bytesneeded:
dq 0
done:
db 0

createThreadAndWait(enumddp1)

[disable]
dealloc(enumddp1)
dealloc(bytesneeded)
dealloc(done)
]]

  r,info=autoAssemble(s1,true)

  if r then
    --wait till done
    --while readBytesLocal(info.allocs.done.address,1)==0 do
    --  sleep(10)
    --end

    local bytesNeeded=readQwordLocal(info.allocs.bytesneeded.address)

    autoAssemble(s1,true,info) --dealloc

    local s2=[[
[enable]
alloc(enumddp2,512)
alloc(driverlist,%d)
alloc(bytesneeded,8)
alloc(done,1)
alloc(result,4)

enumddp2:
sub rsp,28
mov rcx,driverlist
mov rdx,%x
mov r8,bytesneeded
call k32EnumDeviceDrivers
add rsp,28

mov dword [result],eax

mov byte [done],1
ret

bytesneeded:
dq 0
done:
db 0

createThreadAndWait(enumddp2)

[disable]
dealloc(enumddp2)
dealloc(driverlist)
dealloc(bytesneeded)
dealloc(done)
dealloc(result)
  ]]

    local r,info
    r,info=autoAssemble(string.format(s2,bytesNeeded*2, bytesNeeded*2),true)
    if r then
      --while readBytesLocal(info.allocs.done.address,1)==0 do
      --  sleep(10)
      --end

      local result=readIntegerLocal(info.allocs.result.address)
      local DriverListAddress=info.allocs.driverlist.address
      local Count=readIntegerLocal(info.allocs.bytesneeded.address)/8
      --print(string.format("EnumDeviceDrivers returned %d", result))
      --print(string.format("Driverlist at %x", info.allocs.driverlist.address))
      --print(string.format("Count=%d",Count))

      local result={}
      local i
      for i=1,Count do
        local address=readQwordLocal(info.allocs.driverlist.address+(i-1)*8)
        local name=getDriverName(address)
        result[name]={}
        result[name].Address=address
        result[name].Path=getDriverPath(address)
      end


      autoAssemble(s2,true,info)

      return result

    else
      error("getDriverList error 2:"..info)
    end
  else
    error("getDriverList error 1:"..info)
  end
end

function getExports(filename, loadedAddress)
  filename=string.lower(filename) --windows isn't case sensitive
  filename=filename:gsub([[\systemroot\]],[[c:\windows\]])
  filename=filename:gsub([[\%?%?\]],'')

  local stream=createMemoryStream()
  stream.loadFromFile(filename)


  if (byteTableToString(stream.read(2))~='MZ') then
    stream.destroy()
    error('Not a valid executable')
  end

  stream.Position=60;
  local lfanew=stream.readDword();
  stream.Position=lfanew;

  if (byteTableToString(stream.read(2))~='PE') then
    stream.destroy()
    error('Not a valid windows executable')
  end

  stream.position=stream.Position+2
  local Machine=stream.readWord()
  local SectionCount=stream.readWord()
  stream.Position=stream.Position+12
  local OptHeaderSize=stream.readWord()
  stream.position=stream.Position+2

  local OptHeaderStart=stream.Position
  local Magic=stream.readWord()
  local MajorLinkerVersion=stream.readByte()
  local MinorLinkerVersion=stream.readByte()

  local SizeOfCode=stream.readDword()
  stream.Position=stream.Position+8
  local EntryPoint=stream.readDword()
  local BaseOfCode=stream.readDword()

  local BaseOfData
  local ImageBase

  if Machine==0x8664 then
    BaseOfData=nil
    ImageBase=stream.readQword()
    stream.Position=stream.Position+4+4+2+2+2+2+2+2+4+4+4+4+2+2+8+8+8+8
  else
    BaseOfData=stream.readDword()
    ImageBase=stream.readDword()
    stream.Position=stream.Position+4+4+2+2+2+2+2+2+4+4+4+4+2+2+4+4+4+4
  end

  local LoaderFlags=stream.readDword()
  local RVACount=stream.readDword()

  if RVACount~=16 then
    stream.destroy()
    error('This type of module is currently not supported')
  end

  --DataDirectory follows
  local DataDirPosition=stream.Position
  local ImageSectionHeaderPosition=OptHeaderStart+OptHeaderSize

  --parse the sections so VirtualToFile can function
  local sections={}
  stream.Position=ImageSectionHeaderPosition

  local i
  for i=1,SectionCount do
    sections[i]={}
    sections[i].name=byteTableToString(stream.read(8))
    sections[i].misc=stream.readDword()
    sections[i].virtualAddress=stream.readDword()
    sections[i].sizeOfRawData=stream.readDword()
    sections[i].PointerToRawData=stream.readDword()
    stream.position=stream.position+12;
    sections[i].Characteristics=stream.readDword()
    sections[i].Executable=sections[i].Characteristics & IMAGE_SCN_CNT_CODE == IMAGE_SCN_CNT_CODE

  end

  local function VirtualToFile(VA)
    --scan the address in the sections list
    local i
    local offset
    for i=1,#sections do
      if (VA>=sections[i].virtualAddress) and (VA<sections[i].virtualAddress+sections[i].sizeOfRawData) then
        return sections[i].PointerToRawData+(VA-sections[i].virtualAddress)
      end
    end
  end

  local ExportTableRVAInfo=DataDirPosition+0*8

  stream.Position=ExportTableRVAInfo
  local ExportTableAddress=stream.readDword()
  local ExportTableSize=stream.readDword()
  local ExportTablePosition=VirtualToFile(ExportTableAddress)


  local exports={}
  if ExportTablePosition then
    stream.Position=ExportTablePosition+0x10
    local OrdinalBase=stream.readDword()
    local NumberOfFunctions=stream.readDword()
    local NumberOfNames=stream.readDword()
    local OrdinalCount=NumberOfFunctions-NumberOfNames
    local AddressOfFunctions=VirtualToFile(stream.readDword())
    local AddressOfNames=VirtualToFile(stream.readDword())

    for i=1,OrdinalCount do
      local name='Ordinal'..(OrdinalCount+i)
      exports[name]=loadedAddress+readIntegerLocal(stream.Memory+AddressOfFunctions+i*4)
    end

    for i=1,NumberOfNames do
      local name=readStringLocal(stream.Memory+VirtualToFile(readIntegerLocal(stream.Memory+AddressOfNames+i*4)))
      if name then
        exports[name]=loadedAddress+readIntegerLocal(stream.Memory+AddressOfFunctions+(OrdinalCount+i)*4)
      end
    end
  end

  stream.destroy()

  return exports
end