#include "../main.h"
#include "ModuleBase.hpp"

#include <windows.h>
#include "../wine_winternl.h"
#include <detours.h>

#define STATUS_SHARING_VIOLATION 0xC0000043
#define STATUS_ACCESS_DENIED     0xC0000022



class RemoteLockBreak : public ModuleBase {
    inline static decltype(&NtOpenFile) TrueNtOpenFile;
    inline static decltype(&NtCreateFile) TrueNtCreateFile;

    static
    NTSTATUS NTAPI DetourNtOpenFile(OUT PHANDLE FileHandle,
        IN ACCESS_MASK DesiredAccess, IN POBJECT_ATTRIBUTES ObjectAttributes, OUT PIO_STATUS_BLOCK IoStatusBlock,
        IN ULONG ShareAccess, IN ULONG OpenOptions) {
        DWORD iStatus = TrueNtOpenFile(FileHandle, DesiredAccess, ObjectAttributes,
            IoStatusBlock, ShareAccess, OpenOptions);

        if (iStatus == STATUS_SHARING_VIOLATION || iStatus == STATUS_ACCESS_DENIED) {
            if (CloseHandle(FileHandle)) {
                // try operation again now that file is closed
                iStatus = TrueNtOpenFile(FileHandle, DesiredAccess, ObjectAttributes,
                                         IoStatusBlock, ShareAccess, OpenOptions);
            }
        }

        return iStatus;
    }

    static
    NTSTATUS NTAPI DetourNtCreateFile(OUT PHANDLE FileHandle, IN ACCESS_MASK DesiredAccess,
        IN POBJECT_ATTRIBUTES ObjectAttributes, OUT PIO_STATUS_BLOCK IoStatusBlock,
        IN PLARGE_INTEGER AllocationSize OPTIONAL, IN ULONG FileAttributes, IN ULONG ShareAccess,
        IN ULONG CreateDisposition, IN ULONG CreateOptions, IN PVOID EaBuffer OPTIONAL, IN ULONG EaLength) {
        NTSTATUS iStatus = TrueNtCreateFile(FileHandle, DesiredAccess, ObjectAttributes, IoStatusBlock, AllocationSize,
            FileAttributes, ShareAccess, CreateDisposition, CreateOptions, EaBuffer, EaLength);

        if (iStatus == STATUS_SHARING_VIOLATION || iStatus == STATUS_ACCESS_DENIED) {
            if (CloseHandle(FileHandle)) {
                // try operation again now that file is closed
                iStatus = TrueNtCreateFile(FileHandle, DesiredAccess, ObjectAttributes, IoStatusBlock, AllocationSize,
                                           FileAttributes, ShareAccess, CreateDisposition, CreateOptions, EaBuffer, EaLength);
            }
        }

        return iStatus;
    }

public:
    RemoteLockBreak() {
        TrueNtOpenFile = reinterpret_cast<decltype(&NtOpenFile)>(GetProcAddress(LoadLibraryW(L"ntdll.dll"), "NtOpenFile"));
        TrueNtCreateFile = reinterpret_cast<decltype(&NtCreateFile)>(GetProcAddress(LoadLibraryW(L"ntdll.dll"), "NtCreateFile"));
        DetourAttach(&reinterpret_cast<PVOID&>(TrueNtOpenFile), reinterpret_cast<void*>(DetourNtOpenFile));
        DetourAttach(&reinterpret_cast<PVOID&>(TrueNtCreateFile), reinterpret_cast<void*>(DetourNtCreateFile));
    }
    ~RemoteLockBreak() {
        DetourDetach(&reinterpret_cast<PVOID&>(TrueNtOpenFile), reinterpret_cast<void*>(DetourNtOpenFile));
        DetourDetach(&reinterpret_cast<PVOID&>(TrueNtCreateFile), reinterpret_cast<void*>(DetourNtCreateFile));
    }
};