// Author:   Brandon LaCombe
// Date:     February 3, 2006
// License:  Public Domain
#include "CTLSRipper.h"
#include "..\..\FileTools.h"
#include "..\..\remem.h"

// code start

// Class constructor.
CTLSRipper::CTLSRipper()
{
    m_bOutput = FALSE;
    m_bStrip = FALSE;
}

// Exports TLS table.
VOID CTLSRipper::Export(PVOID pvOutput)
{
    if(pvOutput && m_bOutput && m_bStrip == FALSE)
        *PIMAGE_TLS_DIRECTORY(pvOutput) = m_TLS;
}

// Returns size of exported table.
DWORD CTLSRipper::GetSize()
{
    DWORD dwSize;

    dwSize = 0;
    if(m_bOutput && m_bStrip == FALSE)
        dwSize = sizeof(IMAGE_TLS_DIRECTORY);
    return dwSize;
}

// Loads input file. Returns FALSE if no TLS section exists.
BOOL CTLSRipper::LoadFile(PBYTE pbFile)
{
    BOOL              bRet;
    DWORD             dwTLSRva;
    PIMAGE_NT_HEADERS pNt;

    bRet = FALSE;
    m_pbFile = pbFile;
    pNt = PIMAGE_NT_HEADERS(pbFile + PIMAGE_DOS_HEADER(pbFile)->e_lfanew);
    dwTLSRva = pNt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_TLS].VirtualAddress;
    m_pTLS = (PIMAGE_TLS_DIRECTORY)RvaToPointer(pbFile, dwTLSRva);
    if(m_pTLS) bRet = TRUE;
    return bRet;
}

// Rips TLS table from input file.
VOID CTLSRipper::Rip(PVOID pvFile)
{
    m_bOutput = FALSE;
    if(LoadFile((PBYTE)pvFile))
    {
        m_TLS = *m_pTLS;
        if(m_TLS.AddressOfCallBacks)
            m_TLS.AddressOfCallBacks = m_dwNullDwordVa;
        ZeroMemory(m_pTLS, sizeof(IMAGE_TLS_DIRECTORY));
        m_bOutput = TRUE;
    }
}

// Sets the virtual address of a null dword. Used as the callback array.
VOID CTLSRipper::SetNullDwordVa(DWORD dwNullDwordVa)
{
    m_dwNullDwordVa = dwNullDwordVa;
}

// Sets strip flag.
VOID CTLSRipper::Strip(BOOL bStrip)
{
    m_bStrip = bStrip;
}