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

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



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 (NtClose(FileHandle)) {
                // try operation again now that file is closed
                return 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 (NtClose(FileHandle)) {
                // try operation again now that file is closed
                return 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));
    }
};