/**
 * Utilities to disassemble some X86/ X86_64 byte code.
 */
module disassembler;

import beaengine;
import std.format;
import std.algorithm;

private string _prefix;
private string _suffix;
private string _addrFmtSpec;
private string _eol;
private EolMode _eolMode;

static this()
{
    addressSuffix = "h";
    eolMode = EolMode.sys;
}

/// end-of-line modes
enum EolMode
{
    cr, lf, crlf, sys
}

/**
 * Defines the end-of-line used to format the disassembling.
 * Useful under Windows when prettyDisasm() is written to stdout.
 */
@property EolMode eolMode(){return _eolMode;}
/// ditto
@property void eolMode(EolMode mode)
{
    _eolMode = mode;
    import std.ascii : newline;
    with(EolMode) final switch(mode)
    {
        case cr:    _eol = "\r"; break;
        case lf:    _eol = "\n"; break;
        case crlf:  _eol = "\r\n"; break;
        case sys:   _eol = newline; break;
    }
}

/// defines the prefix used to format an address.
@property string addressPrefix(){return _prefix;}
/// ditto
@property void addressPrefix(string p)
{
    _prefix = p;
    updateAddressFormatSpecifier;
}


/// defines the suffix used to format an address.
@property string addressSuffix(){return _suffix;}
/// ditto
@property void addressSuffix(string s)
{
    _suffix = s;
    updateAddressFormatSpecifier;
}


private void updateAddressFormatSpecifier()
{
    if (_prefix.length + _suffix.length > 6)
    {
        _prefix = "";
        _suffix = "h";
        throw new Exception("too long prefix and suffix, default value reset");
    }
    static if (size_t.sizeof == 4)
        _addrFmtSpec = _prefix ~ "%.8X" ~ _suffix;
    else static if (size_t.sizeof == 8)
        _addrFmtSpec = _prefix ~ "%.16X" ~ _suffix;
    else static assert(0, "unsupported pointer size");
}


/**
 * Utility used to format an address
 */
private struct Address
{
    void* _addr;
    alias _addr this;
    string toString() const
    {
        return format(_addrFmtSpec, _addr);
    }
}

/// Array of Disasm. Usually used to store a full function.
alias Sub = DisasmParams* [];

/// Associates an array of Disasm to a particular address.
alias Subs = Sub[void*];

/// Associates an address to an array of address. Used to store the cross references.
alias SubsCrossRefs = void*[][void*];

/// Handles the translation of the disassembler addresses to a symbol.
struct symbolTable
{
    private static string[const(void*)] _symbols;
    private static string* _lastSymbol;
    private static void* _lastAddress;
    private import std.traits: isImplicitlyConvertible, isSomeFunction, PointerTarget;
    private import std.meta: Alias;

    /// Toggles on or off address translation. Works only for CALL and JMP
    static bool enable;

    /// Associates the string symbol to address.
    static void add(bool addType = false, T)(string symbol, const T address) @safe nothrow
    if (isImplicitlyConvertible!(T,void*) || is(T==delegate) || is(PointerTarget!T==function))
    {
        enable = true;
        static if (isImplicitlyConvertible!(T,void*) || is(T==function))
        {
            _symbols[address] = symbol;
            static if (addType)
                _symbols[address] ~= " (" ~ PointerTarget!T.stringof ~ ")";
        }
        else static if (is(T==delegate))
        {
            _symbols[address.funcptr] = symbol;
            static if (addType)
                _symbols[address.funcptr] ~= " (" ~ PointerTarget!T.stringof ~ ")";
        }
        else static assert(0, "unsupported argument type in "
            ~ __PRETTY_FUNCTION__ ~ " : " ~ T.stringof);
    }

    /// Adds a free or a static function to the symbol table
    static void add(alias symbol, string name = "", bool addType = false)() @safe nothrow
    if (is(typeof(symbol)) && isSomeFunction!(typeof(symbol)))
    {
        enable = true;
        auto funPtr = &symbol;
        static if (name.length)
            _symbols[funPtr] = name;
        else
        {
            alias P = Alias!(__traits(parent, symbol));
            static if (is(P == struct) || is(P == union) || is(P == class))
                _symbols[funPtr] = __traits(parent, symbol).stringof ~ '.' ~ __traits(identifier, symbol);
            else
                _symbols[funPtr] = __traits(identifier, symbol);
            static if (addType)
                _symbols[funPtr] ~= " (" ~ PointerTarget!(typeof(funPtr)).stringof ~ ")";
        }
    }

    /// Removes the symbol matching to address.
    static void remove(const void* address) @safe nothrow
    {
        _symbols.remove(address);
    }

    /// Clears the internal container.
    static void clear() nothrow
    {
        foreach(k; _symbols.byKey)
            _symbols.remove(k);
    }

    /// Indicates if a symbol is stored for address.
    deprecated ("use the in operator instead")
    static bool canFind(const void* address) nothrow
    {
        _lastAddress = cast(void*) address;
        return (_lastSymbol = (_lastAddress in _symbols)) != null;
    }

    /// Returns a pointer to the symbol name if address is stored.
    static const(string)* opBinaryRight(string op : "in")(const void* address)
    nothrow
    {
        _lastAddress = cast(void*) address;
        return _lastAddress in _symbols;
    }

    /// Returns the symbol associated to address.
    deprecated ("use the in operator instead")
    static string symbol(const void* address) nothrow
    {
        if (address != _lastAddress)
            canFind(address);
        if (_lastSymbol)
            return *_lastSymbol;
        else
            return "";
    }

    /**
     * Scans an entire module and adds its functions to the table.
     */
    static void addModule(alias mod)() @safe nothrow
    {
        import std.traits: isSomeFunction;
        import std.algorithm: canFind;
        enable = true;
        foreach(memb;__traits(allMembers,mod))
        static if (is(typeof(__traits(getMember, mod, memb))))
            static if (isSomeFunction! (__traits(getMember,mod,memb) ))
                foreach(ov;__traits(getOverloads,mod,memb))
                    static if (canFind(["package","public","export"],__traits(getProtection,ov)))
                        _symbols[&__traits(getMember,mod,memb)] = memb;
    }
}

nothrow unittest
{

    static void foo(uint a){}

    symbolTable.add("a", cast(void*) 0xF);
    symbolTable.add("b", cast(void*) 0xFF);
    symbolTable.add("c", cast(void*) 0xFFF);
    symbolTable.add!(foo, "", true);

    assert((cast(void*) 0xF) in symbolTable);
    assert((cast(void*) 0xFF) in symbolTable);
    assert((cast(void*) 0xFFF) in symbolTable);
    assert((cast(void*) &foo) in symbolTable);

    assert((cast(void*) 0xE) !in symbolTable);
    assert((cast(void*) 0xEE) !in symbolTable);
    assert((cast(void*) 0xEEE) !in symbolTable);
    assert((cast(void*) 0x0) !in symbolTable);

    assert(symbolTable.symbol(cast(void*) 0xF) == "a");
    assert(symbolTable.symbol(cast(void*) 0xFF) == "b");
    assert(symbolTable.symbol(cast(void*) 0xFFF) == "c");

    assert(symbolTable.symbol(cast(void*) 0xEDD) != "c");
    assert(symbolTable.symbol(cast(void*) 0xEDD) == "");

    assert(symbolTable.symbol(cast(void*) &foo) == "foo (void(uint))");

    symbolTable.clear;
    assert(symbolTable._symbols.length == 0);
}


/**
 * Formats an array of Disasm as a string.
 *
 * Params:
 * sub = an array of Disasm.
 * crossRefs = an array of address from where the first sub instruction is called.
 *
 * Returns:
 * a readable representation of the the instructions contained in sub.
 */
string formatSub(const ref Sub sub, const ref void*[]* crossRefs)
{
    static immutable lastLine = ";--------------------------------------------";
    import std.array: Appender;
    import std..string: fromStringz;
    Appender!string result;
    // (numLine + 2 separators) * (address: 2*size_t*sizeof + instruction: 24 chars)
    result.reserve((sub.length + 2) * (size_t.sizeof * 2 + 24));

    result ~= format(";------- SUB " ~ _addrFmtSpec ~ " -------" ~ _eol, sub[0].eip);
    const size_t i = result.data.length;
    if (symbolTable.enable) if (const(string)* s = sub[0].eip in symbolTable)
        result ~= "; NAMED: " ~ *s ~ _eol;
    if (crossRefs != null && crossRefs.length)
        result ~= format("; XREFS: %s " ~ _eol, cast(Address[])*crossRefs);

    string line;
    foreach(instr; sub)
    {
        auto instrText = fromStringz(instr.asString.ptr);
        version(X86_64) if ((instr.instruction.opcode & 0xFF) == 0xE8
            && instr.instruction.addrValue != 0)
        {
            import std..string: rightJustify, split;
            auto splt = split(instrText);
            if (splt[1].length < 16)
            {
                auto len = splt[1].length - 8 + 16;
                instrText = format("%s %s", splt[0], rightJustify(splt[1], len, '0'));
            }
        }
        result ~= format(_addrFmtSpec ~ "  %s", instr.eip, instrText);
        if (symbolTable.enable && instr.instruction.addrValue)
        {
            void* ptr = cast(void*) instr.instruction.addrValue;
            if (const(string)* s = ptr in symbolTable)
                result ~= format(" ; (%s)", *s);
        }
        result ~= _eol;
    }
    result ~= lastLine[0 .. i-1];
    result ~= _eol;
    result ~= _eol;
    return result.data;
}


/**
 * Disassembles a function and returns its string representation.
 *
 * Params:
 * eip = the entry point, a pointer to a function.
 * maxNesting = indicates how many consecutive sub fonctions can be disassembled.
 *
 * Returns:
 * a string representing the function(s).
 *
 * Examples:
 * ---
 * import std.stdio;
 * import disassembler;
 *
 * void main(string[] args)
 * {
 *     // prints the code generated for main().
 *     writeln(prettyDisasm(&main, 1));
 * }
 * ---
 */
string prettyDisasm(void* eip, short maxNesting = 1)
{
    Subs subs;
    SubsCrossRefs xrefs;
    string result;
    if (maxNesting <= 0) maxNesting = 1;
    short nesting = cast(short) (maxNesting * -1);
    //
    disasmSub(eip, subs, xrefs, nesting);
    auto sortedSubs = sort(subs.keys);
    foreach(immutable i; 0 .. sortedSubs.length)
    {
        const(void*) subAddr = sortedSubs[i];
        auto x = (subAddr in xrefs);
        result ~= formatSub(subs[subAddr], x);
    }
    return result;
}

/// ditto
string prettyDisasm(T)(T eip, short maxNesting = 1)
if (is(T==delegate))
{
    return prettyDisasm(eip.funcptr, maxNesting);
}


/**
 * Disassembling kernel.
 *
 * Disassembles recursively from an address.
 *
 * Params:
 * eip = entry point of the function to disassemble. This must point to some byte-code.
 * subs = the associative array filled with the functions instructions.
 * xrefs = the associative array filled with the function callers.
 * nesting = must initially match the inverse of the maximum nested function call the kernel will disassemble.
 */
void disasmSub(void * eip, ref Subs subs, ref SubsCrossRefs xrefs, ref short nesting)
{
    void*[] forward_jumps;

    scope(exit) --nesting;
    if (++nesting > 0)
        return;

    auto loc = eip;
    Sub sub = new Sub(0);
    scope(success) subs[eip] = sub;

    while(true)
    {
        DisasmParams * cur = new DisasmParams;

        // info for the disassembling
        cur.eip = loc;
        version(X86) cur.archi = Archi.ia32;
        version(X86_64) cur.archi = Archi.intel64;

        // disassembles
        const int len = disassemble(cur);
        if (len <= SpecialInfo.OUT_OF_BLOCK)
            break;

        // CALL: disassembles any valid target, stores the cross-reference
        if (cur.instruction.opcode == 0xE8 && cur.instruction.addrValue != 0)
        {
            auto new_eip = cast(void*) cur.instruction.addrValue;
            if (!(new_eip in subs)) disasmSub(new_eip, subs, xrefs, nesting);
            auto refList = new_eip in xrefs;
            if (!refList) xrefs[new_eip] ~= loc;
            else if (!canFind(*refList, loc)) xrefs[new_eip] ~= loc;
        }

        // JX/JNX/JMP, stores destination, it may be located after a RET (still in the same SUB)
        if (cur.instruction.category == InstrCat.CONTROL_TRANSFER
            && cur.instruction.branch != BranchType.RET
            && cur.instruction.branch != BranchType.CALL
            && cur.instruction.addrValue != 0)
        {
            void* forward_loc = cast(void*) cur.instruction.addrValue;
            if (forward_loc > loc && !canFind(forward_jumps, forward_loc))
                forward_jumps ~= forward_loc;
        }

        sub ~= cur;

        // check INT3 after JMP, "tail CALL", real RET is in the destination
        if (sub.length && cur.instruction.opcode == 0xCC
           && sub[$-1].instruction.category == InstrCat.CONTROL_TRANSFER)
        {
            if (sub[$-1].instruction.addrValue != 0)
            {
                auto new_eip = cast(void*) sub[$-1].instruction.addrValue;
                if (!(new_eip in subs)) disasmSub(new_eip, subs, xrefs, nesting);
                auto refList = new_eip in xrefs;
                if (!refList) xrefs[new_eip] ~= loc;
                else if (!canFind(*refList, loc)) xrefs[new_eip] ~= loc;
            }
            break;
        }

        // RET: end of SUB if no more "forward" location.
        if (cur.instruction.category == InstrCat.CONTROL_TRANSFER
            && cur.instruction.branch == BranchType.RET
            && forward_jumps.length == 0) break;

        loc += len;

        // removes a forward reference if location is reached
        auto i = countUntil(forward_jumps, loc);
        if (i != -1)
            forward_jumps = remove(forward_jumps, i);
    }
}

/// ditto
void disasmSub(T)(T eip, ref Subs subs, ref SubsCrossRefs xrefs, ref short nesting)
if (is(T==delegate))
{
    disasmSub(eip.funcptr, subs, xrefs, nesting);
}