#pragma once

struct CookieQuery
{
    duint addr = 0;
    duint ret = 0;
    duint cookieptr = 0;
    ULONG cookie = 0;
    bool removeAddrBp = false;
    bool removeRetBp = false;

    void HandleNtdllLoad(bool isAttached)
    {
        *this = CookieQuery();
        ULONG returnSize = 0;
        NTSTATUS status = NtQueryInformationProcess(fdProcessInfo->hProcess, ProcessCookie, &cookie, sizeof(cookie), &returnSize);
        if(!NT_SUCCESS(status))
        {
            if(!isAttached && valfromstring("ntdll.dll:NtQueryInformationProcess", &addr))
            {
                if(!BpGet(addr, BPNORMAL, nullptr, nullptr))
                {
                    if(SetBPX(addr, UE_BREAKPOINT, cbUserBreakpoint))
                        removeAddrBp = true;
                    else
                        addr = 0;
                }
            }
        }
    }

    bool HandleBreakpoint(duint cip)
    {
        if(addr && addr == cip && Exprfunc::argget(1) == ProcessCookie)
        {
            if(valfromstring("[csp]", &ret))
            {
                cookieptr = Exprfunc::argget(2);
                if(!BpGet(ret, BPNORMAL, nullptr, nullptr))
                {
                    if(SetBPX(ret, UE_BREAKPOINT, cbUserBreakpoint))
                        removeRetBp = true;
                    else
                        ret = 0;
                }
                else
                    removeRetBp = false;
            }
            if(removeAddrBp)
            {
                DeleteBPX(addr);
                BpDelete(addr, BPNORMAL);
                return true;
            }
        }
        else if(ret && ret == cip)
        {
            if(!MemRead(cookieptr, &cookie, sizeof(cookie)))
                cookie = 0, cookieptr = 0;
            if(removeRetBp)
            {
                DeleteBPX(ret);
                BpDelete(ret, BPNORMAL);
                return true;
            }
        }
        return false;
    }
};