using System;
using System.Collections.Generic;
using System.Text;
using System.Diagnostics;

// NOTES: 
// 1. Most references are to Intels manuals, mostly volume 2A and 2B.
// 2. See also http://www.spiralspace.com/Depot/Projects/Disassembler/disassembler_ia32.aspx
// 3. http://webster.cs.ucr.edu/AoA/Windows/HTML/ISA.html

// TODO
// 1. Allow offset throughout, so image[0] has non-zero address for fixed address disassembly
// 2. Allow bounds checking/clipping so out of bounds does not crash disassembly
// 3. Allow setting which feature sets are allowed, by name and/or processor
// 4. Add undocumented, weird opcodes with ways to mark them?
// 5. Add extensions that allow more sophisticated code following, such as virtual functions and table based switch statements
// 6. Add flags to set the code following methods desired (direct only, follow bases, try all starts, etc.)
//    See notes at http://www.stanford.edu/~stinson/paper_notes/stat_anal/anal_stripped_bin.txt and related directory
// 7. Test against my IdaPro Python script to count items found in code. Script at end of this file

namespace Lomont.Binary.X86
    {
    /// <summary>
    /// Represent a class that can disassemble Intel x86 and x64 (16, 32, and 64 bit) code.
    /// Default is 64 bit code. Set the instruction set and memory model before
    /// using if defaults are not sufficient.
    /// TODO - make some useful easy to use and advanced methods.
    /// </summary>
    public class X86Disassembler
        {

        // todo - add simple constructors taking only a List<byte>, and that plus an offset.

        /// <summary>
        /// The list of disassembled instructions.
        /// </summary>
        public List<X86Instruction> Instructions { get; private set; }

        /// <summary>
        /// Create a default disassembler for the given context.
        /// </summary>
        public X86Disassembler(X86DecodeContext context)
            {
            // initialize
            this.Instructions = new List<X86Instruction>();
            this.unexploredStartLocations = new Queue<ulong>();
            this.addressesDecoded = new HashSet<ulong>();
            this.UndecodedSplits = new List<X86Instruction>();
            this.DecodedSplits = new List<X86Instruction>();
            this.context = context;
            }

        /// <summary>
        /// Disassemble the binary, creating a list of x86 instructions.
        /// </summary>
        public void Disassemble(ulong startOffset = 0)
            {
            // add initial start location
            unexploredStartLocations.Enqueue(startOffset);

            // decode until all opcodes added to the known instructions 
            DecodeAll();

            // order instruction stream nicely, insert data sections
            CleanInstructionStream();
            }

        /// <summary>
        /// Disassemble the binary, creating a list of x86 instructions.
        /// </summary>
        public void Disassemble(List<ulong> entryPoints)
            {
            // add initial start locations
            foreach (var entryPoint in entryPoints)
                unexploredStartLocations.Enqueue(entryPoint);

            // decode until all opcodes added to the known instructions 
            DecodeAll();

            // order instruction stream nicely, insert data sections
            CleanInstructionStream();
            }


        /// <summary>
        /// Represent a notice of a decoded instruction.
        /// </summary>
        public class DisassemblerEventArgs : EventArgs
            {
            /// <summary>
            /// The disassembled instruction
            /// </summary>
            public X86Instruction Instruction { get; set; }
            }

        /// <summary>
        /// Raised whenever an instruction is decoded.
        /// </summary>
        public event EventHandler<DisassemblerEventArgs> InstructionDecoded;

        #region Implementation

        /// <summary>
        /// Call whenever a legal instruction is decoded
        /// </summary>
        protected void OnInstructionDecoded(X86Instruction decodedInstruction)
            {
            var e = InstructionDecoded;
            if (e != null)
                e(this, new DisassemblerEventArgs() { Instruction = decodedInstruction });
            }

        /// <summary>
        /// Holds start offsets still to check. Used to track control flow splits.
        /// </summary>
        Queue<ulong> unexploredStartLocations;

        /// <summary>
        /// Holds locations decoded, used to prevent 
        /// multiple jumps into the same start being processed.
        /// </summary>
        HashSet<ulong> addressesDecoded;

        /// <summary>
        /// The decoding context.
        /// </summary>
        public X86DecodeContext context { get; set; }

        /// <summary>
        /// Decode all the data starting at each address in the unexploredStartLocations
        /// </summary>
        void DecodeAll()
            {
            X86Instruction instruction;
            X86InstructionDecoder decoder = new X86InstructionDecoder();

            while (unexploredStartLocations.Count > 0)
                {
                // get a starting point
                ulong start = unexploredStartLocations.Dequeue();
                if (addressesDecoded.Contains(start))
                    continue;
                if (start > (ulong)context.Data.Data.Count)
                    continue; // out of bounds - todo - check? likely a DLL call or such

                context.Data.DataIndex = (ulong)start;
                Debug.WriteLine("Start at {0:X8}", start);
                // disassemble along this code path until dead end reached
                bool more = true;
                while (more)
                    {
                    if (addressesDecoded.Contains(context.Data.DataIndex))
                        break; // already did this one
                    // try to get an instruction
                    instruction = decoder.DecodeInstruction(context);
                    if (instruction == null)
                        {
                        StringBuilder sb = new StringBuilder();
                        for (int i = 0; i < 16; ++i)
                            sb.Append(String.Format("{0:X2} ", context.Data.Data[(int)start + i]));
                        context.OnError("Illegal instruction at\n\t0x{0:X8}: {1}", start, sb.ToString());
                        more = false; // todo - next byte and start again?
                        }
                    else
                        {
                        addressesDecoded.Add(instruction.StartAddress); // store as seen
                        Instructions.Add(instruction); // add to list of decoded
                        OnInstructionDecoded(instruction); // send message
                        more = AddSplits(context, instruction); // and see if more on this line
                        }
                    }
                }
            }

        /// <summary>
        /// Instruction flow splits decoded correctly
        /// </summary>
        public List<X86Instruction> DecodedSplits { get; private set; }

        /// <summary>
        /// Instruction flow splits not decoded correctly
        /// </summary>
        public List<X86Instruction> UndecodedSplits { get; private set; }

        /// <summary>
        /// Add any instruction flow splits to the queue.
        /// Does not add next instruction location.
        /// Return true if there is an instruction following, else false.
        /// </summary>
        /// <param name="instruction">The instruction to check for splits</param>
        /// <param name="context">The decoding context</param>
        /// <returns>true if there is an instruction directly after the current one</returns>
        bool AddSplits(X86DecodeContext context, X86Instruction instruction)
            {
            switch (instruction.Definition.Next)
                {
                case X86InstructionDefinition.BranchType.Jump:
                    context.OnWarning("Unfinished branch type {0}", instruction.Definition.Next);
                    // see if any operands have legit values
                    if (instruction.SourceOperand.OperandFormat == X86Operand.Format.Immediate)
                        {
                        var endLocation = ComputeJumpTarget((ulong)instruction.Length + instruction.StartAddress, instruction.SourceOperand);
                        unexploredStartLocations.Enqueue(endLocation);
                        DecodedSplits.Add(instruction);
                        Debug.WriteLine("Split: {0:X8}->{1:X8}", instruction.StartAddress, endLocation);
                        }
                    else if ((instruction.SourceOperand.OperandFormat == X86Operand.Format.IndirectOffset) && (instruction.SourceOperand.Scale == 4))
                        { // often switch statements are of form JMP offset[ECX*4]. Use table lookups to find values in range
                        ulong imageBase = context.ImageBase;
                        var s = instruction.SourceOperand;
                        var tableAddress = (int)((ulong)s.Value - imageBase);
                        // while these are legal, add to checkable locations
                        DecodedSplits.Add(instruction);
                        ulong entry = 0;
                        do
                            {
                            entry = (ulong)(BitUtils.Read4(context.Data.Data, ref tableAddress) - imageBase);
                            unexploredStartLocations.Enqueue(entry); // todo - ok?
                            } while (entry < (ulong)context.Data.Data.Count);

                        }
                    else
                        { // todo - decode others. Switch statements often of form JMP offset[ECX*4], where offset is pointer to table
                        Debug.WriteLine("Jump: {0}", instruction);
                        UndecodedSplits.Add(instruction);
                        }
                    return false;
                case X86InstructionDefinition.BranchType.Normal:
                    return true;
                case X86InstructionDefinition.BranchType.Return:
                    return false;
                case X86InstructionDefinition.BranchType.Split:
                    context.OnWarning("Unfinished branch type {0}", instruction.Definition.Next);
                    // see if any operands have legit values
                    if (instruction.SourceOperand.OperandFormat == X86Operand.Format.Immediate)
                        {
                        var endLocation = ComputeJumpTarget((ulong)instruction.Length + instruction.StartAddress, instruction.SourceOperand);
                        unexploredStartLocations.Enqueue(endLocation);
                        Debug.WriteLine("Split: {0:X8}->{1:X8}", instruction.StartAddress, endLocation);
                        DecodedSplits.Add(instruction);
                        }
                    return true;
                default:
                    context.OnError("Invalid instruction next operation {0}", instruction.Definition.Next);
                    return false;
                }
            }

        /// <summary>
        /// Given a start address, and an operand with type Immediate, 
        /// compute the destination address.
        /// </summary>
        /// <param name="address"></param>
        /// <param name="operand"></param>
        /// <returns></returns>
        ulong ComputeJumpTarget(ulong address, X86Operand operand)
            {
            switch (operand.Size)
                {
                case X86Instruction.DataSize.Byte:
                    return (ulong)((long)address + (sbyte)operand.Value);
                case X86Instruction.DataSize.Word:
                    return (ulong)((long)address + (short)operand.Value);
                case X86Instruction.DataSize.Dword:
                    return (ulong)((long)address + (int)operand.Value);
                case X86Instruction.DataSize.Qword:
                    return (ulong)((long)address + (long)operand.Value);
                default:
                    var str = String.Format("Unknown operand jump target size {0}", operand.Size);
                    context.OnError(str);
                    throw new ArgumentException(str);
                }
            }

        /// <summary>
        /// Clean the instruction list. Reorders into correct order, 
        /// inserts data blocks for undecoded regions.
        /// </summary>
        void CleanInstructionStream()
            {
            Instructions.Sort((a, b) => a.StartAddress.CompareTo(b.StartAddress));
            List<X86Instruction> tempList = Instructions;
            Instructions = new List<X86Instruction>();
            foreach (var inst in tempList)
                {
                // todo - add data (undecoded) sections
                // todo - add warnings on overlapping instructions? or gaps?
                Instructions.Add(inst);
                }
            }

        #endregion

        }
    }


/*
 * Sample script to count instructions found in IdaPRO to compare to this disassembler.
 * Example testing:
 * 
 * Filename              CLD      IDA       CLD inst  IDA inst   IDA total         Coverage
 *                       inst     inst      bytes       bytes      bytes (+data)    (bytes)
 * pedump.exe           18729    22697     53845        65150       110592            83%
 * EAC.exe              172445   586681    625796       2166925     3903488           29%
 * lame_enc.dll          53987    85734     205473       322510      671744           64%
 * twain_32.dll          5687     10148    16708        29260       49152             57%
 * Bioshock2Launcher.exe 18647   180067    53061       20212098    30552064           0.3 %
 * 
#------------------------------------------------------------
# Chris Lomont written IdaPRO Python tools
#------------------------------------------------------------
"""
LomontIdaProScripts.py - testing functions
"""
import idaapi
import idc
import types

def CountInstructions():
    """
    Count number of disassembled instructions in the current file
    """
    Message("\n\nRunning the instruction counter.\n")
    minEA = idaapi.cvar.inf.minEA
    maxEA = idaapi.cvar.inf.maxEA
    instructions = 0
    dataCount = 0
    dataBytes = 0
    unknownCount = 0
    codeBytes = 0
    Message("Searching range %d to %d.\n" % (minEA,maxEA))
    ea = minEA # minimum effective address
    while ea != BADADDR and ea <= maxEA:
        flags = GetFlags(ea) # get info about current EA
        eaNext = idaapi.next_head(ea, maxEA)
        if isCode(flags):
            instructions = instructions + 1
            if eaNext != BADADDR:
                codeBytes = codeBytes + eaNext - ea
            else :
                codeBytes = codeBytes + maxEA - ea                
        if isData(flags):
            dataCount = dataCount + 1
            if eaNext != BADADDR:
                dataBytes = dataBytes + eaNext - ea
            else:
                dataBytes = dataBytes + maxEA - ea                               
        if isUnknown(flags):
            unknownCount = unknownCount + 1
        ea = eaNext
    Message("%d bytes\n%d instructions in %d bytes\n%d data locations in %d bytes\n%d unknown bytes\n"
            % (maxEA-minEA,instructions, codeBytes,dataCount, dataBytes, unknownCount) )
    Message("%d total bytes, %d in data and code\n"%
            (maxEA-minEA,codeBytes+dataBytes))

# end of file
*/