using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.IO;
using Lomont.Binary.FileFormat;
using Lomont.Binary;
using Lomont.Binary.X86;
using Lomont.Utils;

// todo - make reuse the disassembler - avoid parsing opcodes each file, etc.

namespace Lomont.CLD
    {
    /// <summary>
    /// Recursively scan all executables sub to 
    /// a directory, gathering stats.
    /// </summary>
    public class Scanner
        {

        /// <summary>
        /// Each instruction has this as tuples
        /// </summary>
        public class InstructionStat
            {
            /// <summary>
            /// Ordered prefixes in instruction.
            /// </summary>
            public ulong Prefix;

            /// <summary>
            /// Opcode: prefix, 0F, primary, secondary
            /// </summary>
            public byte OpcodeLength;

            /// <summary>
            /// 0 for none, 1 for ModR/M byte only, 2 for SIB, 3 for both
            /// </summary>
            public byte Modifier;

            /// <summary>
            /// Length of immediate data, 0,1,2,4 bytes
            /// </summary>
            public byte ImmediateLength;

            /// <summary>
            /// Length of displacement data, 0,1,2,4 bytes
            /// </summary>
            public byte DisplacementLength;

            /// <summary>
            /// The mnemonic
            /// </summary>
            public X86Mnemonic Mnemonic;

            /// <summary>
            /// Index into instruction tables
            /// </summary>
            public int InstructionIndex;

            /// <summary>
            /// Populate fields from an instruction
            /// </summary>
            /// <param name="instruction">The instruction to gather stats from</param>
            public InstructionStat(X86Instruction instruction)
                {
                var i = instruction;
                var d = i.Definition;
                foreach (byte b in i.Prefixes)
                    Prefix = (Prefix << 8) | b;
                if (d.Prefix.HasValue)
                    OpcodeLength++;
                if (d.OF.HasValue)
                    OpcodeLength++;
                OpcodeLength++;
                if (d.SecondaryOpcode.HasValue)
                    OpcodeLength++;
                if (i.HasModRM)
                    Modifier |= 1;
                if (i.HasSIB)
                    Modifier |= 2;
                ImmediateLength = (byte)i.ImmediateSize;
                DisplacementLength = (byte)i.DisplacementSize;
                Mnemonic = d.Mnenomic;
                InstructionIndex = d.Index;
                }

            #region Equality

            /// <summary>
            /// Return true if two objects are equal
            /// meaning, they have equal fields
            /// </summary>
            /// <param name="obj"></param>
            /// <returns></returns>
            public override bool Equals(object obj)
                {
                var a = this;
                var b = obj as InstructionStat;
                if (b == null)
                    return false;
                return
                    (a.InstructionIndex == b.InstructionIndex) &&
                    (a.DisplacementLength == b.DisplacementLength) &&
                    (a.ImmediateLength == b.ImmediateLength) &&
                    (a.Modifier == b.Modifier) &&
                    (a.OpcodeLength == b.OpcodeLength) &&
                    (a.Prefix == b.Prefix);
                }

            public override int GetHashCode()
                {
                return (int)(
                    Prefix ^ (ulong)(OpcodeLength << 4) ^
                    ((ulong)Modifier << 8) ^ ((ulong)ImmediateLength << 16) ^
                    ((ulong)DisplacementLength) ^ (ulong)(InstructionIndex)
                );
                }
            #endregion

            }

        /// <summary>
        /// the stats to gather for each file
        /// </summary>
        public class FileStats
            {
            /// <summary>
            /// Complete filename including path
            /// </summary>
            public string Filename;
            /// <summary>
            /// bytes in the code section. Also includes exports, imports, other items
            /// </summary>
            public long CodeBytesRead;
            /// <summary>
            /// Number of instructions decoded
            /// </summary>
            public long InstructionsDecoded;
            /// <summary>
            /// Number of bytes in decoded instructions
            /// </summary>
            public long InstructionBytesDecoded;
            /// <summary>
            /// Time to disassemble in ms
            /// </summary>
            public long DisassembleTime;
            /// <summary>
            /// instruction overlaps bytes
            /// </summary>
            public long OverlappedInstructionBytes;
            /// <summary>
            /// Bytes in between decoded instructions not decoded
            /// </summary>
            public long CodeGapBytes;

            /// <summary>
            /// Count how many time each type appears
            /// </summary>
            public Dictionary<InstructionStat, ulong> InstructionStats = new Dictionary<InstructionStat, ulong>();
            }

        /// <summary>
        /// Stats gathered globally
        /// </summary>
        public class TotalStats
            {
            /// <summary>
            /// total file count
            /// </summary>
            public long FilesTouched;
            /// <summary>
            /// total PE file count
            /// </summary>
            public long PEFiles;
            /// <summary>
            /// total bytes of files scanned
            /// </summary>
            public long FileBytes;
            /// <summary>
            /// total bytes of PE files scanned
            /// </summary>
            public long PEFileBytes;

            /// <summary>
            /// Aggregate all files
            /// </summary>
            public FileStats FileTotals = new FileStats();

            /// Start time of entire process
            /// </summary>
            public DateTime StartTime;
            /// <summary>
            /// End time of entire process
            /// </summary>
            public DateTime EndTime;

            }

        /// <summary>
        /// sum global stats here
        /// </summary>
        public TotalStats Totals = new TotalStats();

        /// <summary>
        /// Read a file, see if it is a PE file, and if so, read 
        /// the executable part out, set all start addresses.
        /// If not PE or fails, return false, else return true.
        /// </summary>
        /// <param name="filename">The filename to read</param>
        /// <param name="entryPoints">A list of entry points to start disassembling from</param>
        /// <param name="bitDepth">The bitdepth of the exe (32 or 64)</param>
        /// <param name="data">The image to decode</param>
        /// <param name="imageBase">The image base, which is the desired PE loading location</param>
        /// <param name="stats">Stats about the file</param>
        /// <returns></returns>
        bool ReadPEFile(string filename, out List<ulong> entryPoints, out PEFormat.BitDepths bitDepth, out List<byte> data, out ulong imageBase, FileStats stats)
            {
            entryPoints = new List<ulong>();
            bitDepth = PEFormat.BitDepths.Unknown;
            data = null;
            imageBase = 0;
            try
                {
                PEFormat format = new PEFormat();
                //                format.MessageEvent += (o, e) => Console.WriteLine(e.Message);

                // avoid reading file if too big, else memory problems
                const int cutoff = 10000000; // about 10 MB cutoff for now. If it is a PE file, then entire thing read
                if (new FileInfo(filename).Length >= cutoff)
                    {
                    // test on a small parse, and if passes, continue
                    using (var file = File.OpenRead(filename))
                        {
                        byte[] array = new byte[cutoff + 10];
                        file.Read(array, 0, cutoff);
                        data = array.ToList();
                        }
                    if (!format.ParseSmall(data))
                        return false;
                    }
                data = File.ReadAllBytes(filename).ToList();
                if (data.Count == 0)
                    return false;
                bool success = format.Parse(data);
                stats.Filename = filename;
                if (!success)
                    return false;
                else
                    { // locate data and start
                    bitDepth = format.BitDepth;
                    List<uint> entryPointsInt;
                    format.GetCodeAndEntryPoints(out data, out entryPointsInt);
                    foreach (var entry in entryPointsInt)
                        entryPoints.Add(entry);
                    // get code segment size. The array length might be different
                    stats.CodeBytesRead = format.OptionalHeader.SizeOfCode;
                    imageBase = format.OptionalHeader.ImageBase;
                    return true;
                    }
                }
            catch (Exception ex)
                {
                OnScanEvent(ScanEventArgs.Type.Message, null,
                    String.Format("ERROR: Exception {0} parsing file {1}", ex.Message, filename)
                    );
                return false;
                }
            }

        /// <summary>
        /// Disassemble the binary. Return true if successful, else false
        /// </summary>
        /// <param name="binary">The binary to disassemble</param>
        /// <param name="type">The bit depth</param>
        /// <param name="entryPoints">A list of entry points to decode from</param>
        /// <param name="imageBase">the PE file image base</param>
        /// <param name="stats">Stats to collect</param>
        /// <returns>true on success, else false</returns>
        bool Disassemble(List<byte> binary, PEFormat.BitDepths type, List<ulong> entryPoints, ulong imageBase, FileStats stats)
            {
            ByteStream data = new ByteStream(binary, 0);
            X86ProcessorMode mode = new X86ProcessorMode();
            if (type == PEFormat.BitDepths.Depth32)
                mode.OperationMode = X86ProcessorMode.OperationModes.Protected;
            else if (type == PEFormat.BitDepths.Depth64)
                mode.OperationMode = X86ProcessorMode.OperationModes.SixtyFourBit;
            else
                {
                OnScanEvent(ScanEventArgs.Type.Message, null, "ERROR: unknown image bit depth");
                return false;
                }
            Messager messager = new Messager();
            //messager.MessageEvent += (o, e) => Console.WriteLine(e.Message);
            X86DecodeContext context = new X86DecodeContext(messager, mode, data);
            context.ImageBase = imageBase;
            X86Disassembler dis = new X86Disassembler(context);

            //dis.InstructionDecoded += (o, e) => Console.WriteLine("Decoded: {0:X8}: {1}", e.Instruction.StartAddress, e.Instruction);

            long startTime = Environment.TickCount;
            dis.Disassemble(entryPoints);
            long endTime = Environment.TickCount;

            // record stats
            stats.DisassembleTime = endTime - startTime;
            stats.InstructionsDecoded = dis.Instructions.Count;
            stats.InstructionBytesDecoded = dis.Instructions.Sum(s => s.Length);

            // process disassembled instructions
            int max = dis.Instructions.Count;
            stats.CodeGapBytes = 0;
            stats.OverlappedInstructionBytes = 0;
            for (int i = 0; i < max; ++i)
                {
                // count overlaps
                if (i < max - 1)
                    {
                    // step between neighboring instructions
                    int delta = (int)(dis.Instructions[i + 1].StartAddress - dis.Instructions[i].StartAddress);
                    delta -= (int)dis.Instructions[i].Length;
                    if (delta > 0)
                        stats.CodeGapBytes += delta;
                    if (delta < 0)
                        stats.OverlappedInstructionBytes += -delta;
                    }

                var ins = dis.Instructions[i];
                var insStat = new InstructionStat(ins);
                if (!stats.InstructionStats.ContainsKey(insStat))
                    stats.InstructionStats.Add(insStat, 0);
                stats.InstructionStats[insStat]++;
                }
            return true;
            } // Disassemble

        /// <summary>
        /// Scan the files recursively from here, tag those as PE files that are
        /// </summary>
        /// <param name="fileroot"></param>
        public void Scan(string fileroot)
            {
            Totals.StartTime = DateTime.Now;
            OnScanEvent(ScanEventArgs.Type.Start, null, null);
            Recurse(fileroot);
            OnScanEvent(ScanEventArgs.Type.End, null, null);
            }

        /// <summary>
        /// Scanner event args, send messages to higher level listeners
        /// </summary>
        public class ScanEventArgs : EventArgs
            {
            public enum Type
                {
                /// <summary>
                /// Scanner starting
                /// </summary>
                Start,
                /// <summary>
                /// Beginning to try file
                /// </summary>
                FileBegin,
                /// <summary>
                /// File processed
                /// </summary>
                FileDone,
                /// <summary>
                /// File failed, and reason.
                /// </summary>
                FileFailed,
                /// <summary>
                /// New directory
                /// </summary>
                Directory,
                /// <summary>
                /// Show message. Should start with INFO, WARNING, or ERROR
                /// </summary>
                Message,
                /// <summary>
                /// Scanner done
                /// </summary>
                End
                }
            /// <summary>
            /// Reason for this event
            /// </summary>
            public Type Reason;
            /// <summary>
            /// File stats if one found
            /// </summary>
            public FileStats Stats;
            /// <summary>
            /// A text message
            /// </summary>
            public string Message;
            /// <summary>
            /// This filename was processed, but not a valid PE file
            /// </summary>
            public string Filename;
            }

        /// <summary>
        /// Event triggered for scan events
        /// </summary>
        public event EventHandler<ScanEventArgs> ScanEvent;

        /// <summary>
        /// Trigger a scanning event
        /// </summary>
        protected void OnScanEvent(ScanEventArgs.Type reason, FileStats fileStats, string Message)
            {
            var e = ScanEvent;
            if (e != null)
                e(this, new ScanEventArgs() { Message = Message, Reason = reason, Stats = fileStats });
            }

        /// <summary>
        /// Trigger a scanning event, used for flagging a failed filename
        /// </summary>
        protected void OnScanEvent(ScanEventArgs.Type reason, string filename)
            {
            var e = ScanEvent;
            if (e != null)
                e(this, new ScanEventArgs() { Reason = reason, Filename = filename });
            }

        /// <summary>
        /// process a single filename. If PE file, disassemble, else just
        /// count it
        /// </summary>
        /// <param name="filename">The filename to process</param>
        void ProcessFile(string filename)
            {
            try
                {
                OnScanEvent(ScanEventArgs.Type.FileBegin, filename);

                PEFormat.BitDepths bitDepth;
                List<byte> data;
                List<ulong> entryPoints;
                var stats = new FileStats();
                ulong imageBase;
                bool success = ReadPEFile(filename, out entryPoints, out bitDepth, out data, out imageBase, stats);
                if (success)
                    { // obtained valid PE file, process it
                    Disassemble(data, bitDepth, entryPoints, imageBase, stats);
                    AddStats(stats, data.Count);
                    Totals.PEFiles++;
                    }
                else
                    {
                    if (data != null)
                        Totals.FileBytes += data.Count;
                    Totals.FilesTouched++;
                    OnScanEvent(ScanEventArgs.Type.FileFailed, null, "not PE");
                    }
                }
            catch (Exception ex)
                {
                OnScanEvent(ScanEventArgs.Type.Message, null,
                    String.Format("ERROR: File scanning exception {0}", ex.Message)
                    );
                }
            }

        /// <summary>
        /// Given the stats from a decoded file, process them
        /// </summary>
        /// <param name="stats">The per file stats</param>
        /// <param name="fileLength">The length in bytes of the file</param>
        void AddStats(FileStats stats, int fileLength)
            {
            var totalStats = Totals.FileTotals;
            totalStats.CodeBytesRead += stats.CodeBytesRead;
            totalStats.CodeGapBytes += stats.CodeGapBytes;
            totalStats.DisassembleTime += stats.DisassembleTime;
            totalStats.Filename = "";
            totalStats.InstructionBytesDecoded += stats.InstructionBytesDecoded;
            totalStats.InstructionsDecoded += stats.InstructionsDecoded;
            totalStats.OverlappedInstructionBytes += stats.OverlappedInstructionBytes;


            // aggregate instruction stats
            var ti = totalStats.InstructionStats;
            foreach (var stat in stats.InstructionStats)
                {
                if (!ti.ContainsKey(stat.Key))
                    ti.Add(stat.Key, stat.Value);
                else
                    ti[stat.Key] += stat.Value;

                }

            Totals.FileBytes += fileLength;
            Totals.FilesTouched++;
            Totals.PEFileBytes += fileLength;

            // mark event
            OnScanEvent(ScanEventArgs.Type.FileDone, stats, null);
            }

        /// <summary>
        /// Recursively scan all files under path.
        /// </summary>
        /// <param name="path">The path to start with</param>
        void Recurse(string path)
            {
            try
                {
                OnScanEvent(ScanEventArgs.Type.Directory, path);
                var files = Directory.GetFiles(path);
                foreach (var file in files)
                    ProcessFile(file);
                var dirs = Directory.GetDirectories(path);
                foreach (var dir in dirs)
                    Recurse(dir);
                }
            catch (Exception ex)
                {
                OnScanEvent(ScanEventArgs.Type.Message, null,
                    String.Format("ERROR: Recurse exception {0}", ex.Message)
                    );
                }
            }
        }
    }
