import argparse
import ctypes
import os.path
from pathlib import Path

import win32file


class FILE_DIRECTORY_INFORMATION(ctypes.Structure):
    _fields_ = [
        ("NextEntryOffset", ctypes.c_uint32),
        ("FileIndex", ctypes.c_uint32),
        ("CreationTime", ctypes.c_int64),
        ("LastAccessTime", ctypes.c_int64),
        ("LastWriteTime", ctypes.c_int64),
        ("ChangeTime", ctypes.c_int64),
        ("EndOfFile", ctypes.c_int64),
        ("AllocationSize", ctypes.c_int64),
        ("FileAttributes", ctypes.c_uint32),
        ("FileNameLength", ctypes.c_uint32),
        ("FileName", ctypes.c_wchar),
    ]


def main():
    parser = argparse.ArgumentParser("list directory using NtQueryOpenFile")
    parser.add_argument("-r", "--recurse", action="store_true", help="list recursively")
    parser.add_argument("directory", type=Path, help="directory to list")

    args = parser.parse_args()
    recursive: bool = args.recurse
    directory: Path = args.directory

    kernel32 = ctypes.windll.LoadLibrary("kernel32.dll")
    ntdll = ctypes.windll.LoadLibrary("ntdll.dll")

    h = kernel32.CreateFileW(
        os.path.expanduser(str(directory)),
        win32file.GENERIC_READ,
        win32file.FILE_SHARE_READ | win32file.FILE_SHARE_WRITE,
        0,
        3,
        win32file.FILE_FLAG_BACKUP_SEMANTICS,
        0,
    )

    if h == -1:
        print("failed to open directory")
        exit(1)

    print(f"handle={h}")

    io_status_block = ctypes.create_string_buffer(64)
    length = 1024
    file_information_buffer = ctypes.create_string_buffer(length)

    while (
        ntdll.NtQueryDirectoryFile(
            h,
            0,
            0,
            0,
            io_status_block,
            file_information_buffer,
            length,
            1,
            False,
            0,
            False,
        )
        == 0
    ):
        index = 0
        while index < length:
            file_information = FILE_DIRECTORY_INFORMATION.from_buffer(
                file_information_buffer, index
            )
            print(
                "{name} ({length}) -> {offset}".format(
                    offset=file_information.NextEntryOffset,
                    length=file_information.FileNameLength,
                    name=ctypes.wstring_at(
                        ctypes.addressof(file_information)
                        + FILE_DIRECTORY_INFORMATION.FileName.offset,
                        file_information.FileNameLength // 2,
                    ),
                )
            )

            if file_information.NextEntryOffset == 0:
                break

            index += file_information.NextEntryOffset

    kernel32.CloseHandle(h)


if __name__ == "__main__":
    main()