#include "pe/ResourceDirWrapper.h"
#include "pe/PEFile.h"

#define MAX_ENTRIES (-1)
#define MAX_DEPTH 5

/*
typedef struct _IMAGE_RESOURCE_DIRECTORY {
    DWORD   Characteristics;
    DWORD   TimeDateStamp;
    WORD    MajorVersion;
    WORD    MinorVersion;
    WORD    NumberOfNamedEntries;
    WORD    NumberOfIdEntries;
    //  IMAGE_RESOURCE_DIRECTORY_ENTRY DirectoryEntries[];
} IMAGE_RESOURCE_DIRECTORY, *PIMAGE_RESOURCE_DIRECTORY;


 // Each directory contains the 32-bit Name of the entry and an offset,
 // relative to the beginning of the resource directory of the data associated
 // with this directory entry.  If the name of the entry is an actual text
 // string instead of an integer Id, then the high order bit of the name field
 // is set to one and the low order 31-bits are an offset, relative to the
 // beginning of the resource directory of the string, which is of type
 // IMAGE_RESOURCE_DIRECTORY_STRING.  Otherwise the high bit is clear and the
 // low-order 16-bits are the integer Id that identify this resource directory
 // entry. If the directory entry is yet another resource directory (i.e. a
 // subdirectory), then the high order bit of the offset field will be
 // set to indicate this.  Otherwise the high bit is clear and the offset
 // field points to a resource data entry.

#define RESOURCE_NAME_IS_STRING        0x80000000
#define RESOURCE_DATA_IS_DIRECTORY     0x80000000


typedef struct _IMAGE_RESOURCE_DIRECTORY_ENTRY {
    union {
        struct {
            DWORD NameOffset:31;
            DWORD NameIsString:1;
        } name;
        DWORD   Name;
        WORD    Id;
    };
    union {
        DWORD   OffsetToData;
        struct {
            DWORD   OffsetToDirectory:31;
            DWORD   DataIsDirectory:1;
        } dir;
    };
} IMAGE_RESOURCE_DIRECTORY_ENTRY, *PIMAGE_RESOURCE_DIRECTORY_ENTRY;

*/
//-------------

IMAGE_RESOURCE_DIRECTORY* ResourceDirWrapper::mainResourceDir()
{
    offset_t rva = getDirEntryAddress();

    BYTE *ptr = m_Exe->getContentAt(rva, Executable::RVA, sizeof(IMAGE_RESOURCE_DIRECTORY));
    return (IMAGE_RESOURCE_DIRECTORY*) ptr;
}

IMAGE_RESOURCE_DIRECTORY* ResourceDirWrapper::resourceDir()
{
    if (this->rawOff == 0) return mainResourceDir();

    BYTE *ptr = m_Exe->getContentAt(this->rawOff, Executable::RAW, sizeof(IMAGE_RESOURCE_DIRECTORY));
    return (IMAGE_RESOURCE_DIRECTORY*) ptr;
}

bool ResourceDirWrapper::wrap()
{
    clear();
    if (this->topEntryID == TOP_ENTRY_ROOT && this->album != NULL) {
        this->album->clear();
    }
    IMAGE_RESOURCE_DIRECTORY* dir = resourceDir();
    if (dir == NULL) return false;
    size_t namesNum = dir->NumberOfNamedEntries;
    size_t idsNum = dir->NumberOfIdEntries;

    size_t totalEntries = namesNum + idsNum;
    for (size_t i = 0; i < totalEntries && i < MAX_ENTRIES; i++ ) {

        long topDirId = (this->topEntryID != TOP_ENTRY_ROOT) ? this->topEntryID : long(i) ;

        //ResourceEntryWrapper(PEFile *pe, ResourceDirWrapper *parentDir, size_t entryNumber, long topEntryId, ResourcesAlbum *resAlbum)
        ResourceEntryWrapper* entry = new ResourceEntryWrapper(this->m_PE, this, i, topDirId, this->album);

        if (entry->getPtr() == NULL) {
            delete entry;
            break;
        }
        if (this->topEntryID == TOP_ENTRY_ROOT && this->album != NULL) {
            pe::resource_type typeId = static_cast<pe::resource_type>(entry->getID());
            this->album->mapIdToLeafType(i, typeId);
        }
        //this->parsedSize += val;
        this->entries.push_back(entry);
    }
    //printf("Entries: %d\n", getEntriesCount());
    return true;
}


bufsize_t ResourceDirWrapper::getSize()
{
    if (getPtr() == NULL) return 0;
    bufsize_t size = sizeof(IMAGE_RESOURCE_DIRECTORY);
    size += getEntriesAreaSize();
    return size;
}

void* ResourceDirWrapper::getFieldPtr(size_t fId, size_t subField)
{
    IMAGE_RESOURCE_DIRECTORY* d = resourceDir();
    if (d == NULL) return NULL;

    switch (fId) {
        case CHARACTERISTIC: return &d->Characteristics;
        case TIMESTAMP: return &d->TimeDateStamp;
        case MAJOR_VER: return &d->MajorVersion;
        case MINOR_VER: return &d->MinorVersion;
        case NAMED_ENTRIES_NUM : return &d->NumberOfNamedEntries;
        case ID_ENTRIES_NUM : return &d->NumberOfIdEntries;
    }
    return this->getPtr();
}

bufsize_t ResourceDirWrapper::getFieldSize(size_t fId, size_t subField)
{
    IMAGE_RESOURCE_DIRECTORY* d = resourceDir();
    if (d == NULL) return 0;

    // calculate size of the last field separately:
    if (fId == ID_ENTRIES_NUM) {
        void *nextPtr = (&d->NumberOfIdEntries) + 1;
        void *argPtr = &d->NumberOfIdEntries;
        return (BYTE*) nextPtr - (BYTE*)argPtr;
    }
    return DataDirEntryWrapper::getFieldSize(fId, subField);
}

QString ResourceDirWrapper::getFieldName(size_t fieldId)
{
    switch (fieldId) {
        case CHARACTERISTIC: return "Characteristics";
        case TIMESTAMP: {
            PEFile* myPe = dynamic_cast<PEFile*>(this->m_Exe);
            if (myPe && myPe->isReproBuild()) {
                return "ReproChecksum";
            }
            return "TimeDateStamp";
        }
        case MAJOR_VER: return "MajorVersion";
        case MINOR_VER: return "MinorVersion";
        case NAMED_ENTRIES_NUM : return "NumberOfNamedEntries";
        case ID_ENTRIES_NUM : return "NumberOfIdEntries";
    }
    return getName();
}

//-----------------------------------------------------------------

void ResourceEntryWrapper::clear()
{
    delete childDir;
    this->childDir = NULL;
    //---
    delete childLeaf;
    this->childLeaf = NULL;
}

bool ResourceEntryWrapper::wrap()
{
    clear();
    //---
    offset_t childRaw = getChildAddress();
    if (childRaw == 0 || childRaw == INVALID_ADDR) return false;

    if (this->isDir()) {
        long depth = this->parentDir->getDepth() + 1;
        if (depth >= MAX_DEPTH) return false;
        this->childDir = new ResourceDirWrapper(this->m_PE, this->album, childRaw, depth, topEntryID);
    } else {
        this->childLeaf = new ResourceLeafWrapper(m_Exe, childRaw, topEntryID);
        if (this->album != NULL) {
            album->putLeaf(childLeaf, topEntryID);
        }
    }
    return true;
}


IMAGE_RESOURCE_DIRECTORY_ENTRY* ResourceEntryWrapper::getEntryPtr()
{
    if (this->parentDir == NULL) return NULL;
    uint64_t offset = parentDir->getFieldOffset(ResourceDirWrapper::ID_ENTRIES_NUM);
    if (offset == INVALID_ADDR) return NULL;

    size_t lastSize = parentDir->getFieldSize(ResourceDirWrapper::ID_ENTRIES_NUM, FIELD_NONE);
    offset += lastSize;
    offset += (this->entryNum * sizeof(IMAGE_RESOURCE_DIRECTORY_ENTRY));

    void *ptr = m_Exe->getContentAt(offset, Executable::RAW, sizeof(IMAGE_RESOURCE_DIRECTORY_ENTRY));
    return (IMAGE_RESOURCE_DIRECTORY_ENTRY*) ptr;
}

//IMAGE_RESOURCE_DIRECTORY_ENTRY
void* ResourceEntryWrapper::getFieldPtr(size_t fieldId, size_t subField)
{
    IMAGE_RESOURCE_DIRECTORY_ENTRY* entry = getEntryPtr();
    if (entry == NULL) return NULL;

    switch (fieldId) {
        case NAME_ID_ADDR:
            return &entry->Name;
        case OFFSET_TO_DATA:
            return &entry->OffsetToData;
    }
    return getPtr();
}

QString ResourceEntryWrapper::getFieldName(size_t fieldId)
{
    switch (fieldId) {
        case NAME_ID_ADDR:
            return(isByName()) ? "Name": "ID";
        case OFFSET_TO_DATA:
            return "Offset To Data";
    }
    return getName();
}

bool ResourceEntryWrapper::isByName()
{
    IMAGE_RESOURCE_DIRECTORY_ENTRY* entry = getEntryPtr();
    if (entry == NULL) return false;
    if (entry->NameIsString == 1) return true;
    return false;
}

bool ResourceEntryWrapper::isDir()
{
    IMAGE_RESOURCE_DIRECTORY_ENTRY* entry = getEntryPtr();
    if (entry == NULL) return false;
    if (entry->DataIsDirectory == 1) return true;
    return false;
}

offset_t ResourceEntryWrapper::getNameOffset()
{
    IMAGE_RESOURCE_DIRECTORY_ENTRY* entry = getEntryPtr();
    if (entry == NULL) return INVALID_ADDR;

    if (this->parentDir == NULL) return INVALID_ADDR;
    offset_t fOff = this->parentDir->getOffset(this->parentDir->getPtr());
    if (fOff == INVALID_ADDR) return INVALID_ADDR;

    if (entry->NameIsString != 1) return INVALID_ADDR;
    return fOff + entry->NameOffset;
}

IMAGE_RESOURCE_DIRECTORY_STRING* ResourceEntryWrapper::getNameStr()
{
    offset_t nameOff = getNameOffset();
    if (nameOff == INVALID_ADDR) return NULL;

    const size_t BASIC_SIZE = sizeof(IMAGE_RESOURCE_DIRECTORY_STRING);
    IMAGE_RESOURCE_DIRECTORY_STRING *ptr = (IMAGE_RESOURCE_DIRECTORY_STRING*) this->m_Exe->getContentAt(nameOff, Executable::RAW, BASIC_SIZE);
    return ptr;
}

QString ResourceEntryWrapper::translateType(WORD id)
{
    switch (id) {
        case pe::RESTYPE_CURSOR : return "Cursor";
        case pe::RESTYPE_FONT :return "Font";
        case pe::RESTYPE_BITMAP : return "Bitmap";
        case pe::RESTYPE_ICON : return "Icon";
        case pe::RESTYPE_MENU : return "Menu";
        case pe::RESTYPE_DIALOG : return "Dialog";
        case pe::RESTYPE_STRING : return "String";
        case pe::RESTYPE_FONTDIR : return "Font Dir";
        case pe::RESTYPE_ACCELERATOR : return "Accelerator";
        case pe::RESTYPE_RCDATA : return "RC Data";
        case pe::RESTYPE_MESSAGETABLE : return "Message Table";

        case pe::RESTYPE_GROUP_CURSOR : return "Cursors Group";
        case pe::RESTYPE_GROUP_ICON : return "Icons Group";
        case pe::RESTYPE_VERSION : return "Version";
        case pe::RESTYPE_DLGINCLUDE : return "Dlg Include";
        case pe::RESTYPE_PLUGPLAY : return "Plug & Play";
        case pe::RESTYPE_VXD : return "VXD";
        case pe::RESTYPE_ANICURSOR : return "Animated Cursor";
        case pe::RESTYPE_ANIICON : return "Animated Icon";
        case pe::RESTYPE_HTML : return "HTML";
        case pe::RESTYPE_MANIFEST : return "Manifest";
    }
    return "";
}

WORD ResourceEntryWrapper::getID()
{
    IMAGE_RESOURCE_DIRECTORY_ENTRY* entry = getEntryPtr();
    if (entry == NULL) return 0;
    if (this->isByName()) return 0;

    return entry->Id;
}

offset_t ResourceEntryWrapper::getChildOffsetToDirectory()
{
    const IMAGE_RESOURCE_DIRECTORY_ENTRY* entry = getEntryPtr();
    if (!entry) {
        return INVALID_ADDR;
    }
    if (entry->DataIsDirectory) {
        return entry->OffsetToDirectory;
    }
    return entry->OffsetToData;
}

offset_t ResourceEntryWrapper::getChildAddress()
{
    if (!this->parentDir) return INVALID_ADDR;
    
    const offset_t fOff = this->parentDir->getOffset(this->parentDir->mainResourceDir());
    if (fOff == INVALID_ADDR) return INVALID_ADDR;

    const offset_t childOff = getChildOffsetToDirectory();
    if (childOff == INVALID_ADDR) return INVALID_ADDR;

    return fOff + childOff;
}

//-------------------------------------------------------------------------

IMAGE_RESOURCE_DATA_ENTRY* ResourceLeafWrapper::leafEntryPtr()
{
    if (this->offset == 0 || this->offset == INVALID_ADDR) return NULL;

    IMAGE_RESOURCE_DATA_ENTRY* leaf = (IMAGE_RESOURCE_DATA_ENTRY*) this->m_Exe->getContentAt(this->offset, Executable::RAW, sizeof(IMAGE_RESOURCE_DATA_ENTRY));
    return leaf;
}

void* ResourceLeafWrapper::getFieldPtr(size_t fieldId, size_t subField)
{
    IMAGE_RESOURCE_DATA_ENTRY* leaf = leafEntryPtr();
    if (!leaf) return NULL;

    switch (fieldId) {
        case OFFSET_TO_DATA: return &leaf->OffsetToData;
        case DATA_SIZE: return &leaf->Size;
        case CODE_PAGE: return &leaf->CodePage;
        case RESERVED: return &leaf->Reserved;
    }
    return NULL;
}

QString ResourceLeafWrapper::getFieldName(size_t fieldId)
{
    switch (fieldId) {
        case OFFSET_TO_DATA: return "OffsetToData";
        case DATA_SIZE: return "DataSize";
        case CODE_PAGE: return "CodePage";
        case RESERVED: return "Reserved";
    }
    return "";
}