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

namespace Lomont.CLD
    {
    /// <summary>
    /// The main disassembly program
    /// Can disassemble binaries, PE files, and dump PE file items
    /// </summary>
    class Program
        {
        /// <summary>
        /// Program version
        /// </summary>
        static string Version = "0.1";

        /// <summary>
        /// Disassemble the binary
        /// </summary>
        /// <param name="binary"></param>
        static void Disassemble(List<byte> binary, PEFormat.BitDepths type, List<ulong> entryPoints, ulong imageBase)
            {
            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
                {
                Console.WriteLine("ERROR: unknown image bit depth");
                Environment.Exit(-4);
                }
            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);

            dis.Disassemble(entryPoints);
            foreach (var opcode in dis.Instructions)
                Console.WriteLine("{0:X8}: {1}", opcode.StartAddress, opcode);
            Console.WriteLine("{0} instructions decoded in {1} bytes", 
                dis.Instructions.Count,
                dis.Instructions.Sum(ins=>ins.Length)
                );
            Console.WriteLine("{0} decoded splits",dis.DecodedSplits.Count);
            Console.WriteLine("{0} undecoded splits", dis.UndecodedSplits.Count);
            }

        /// <summary>
        /// Disassemble a PE file
        /// </summary>
        /// <param name="PEFile"></param>
        static void DisassemblePE(List<byte> fileImage)
            {
            PEFormat format = new PEFormat();
            format.MessageEvent += (o, e) => Console.WriteLine(e.Message);
            bool success = format.Parse(fileImage);
            if (!success)
                Console.WriteLine("Parsing PE Failed");
            else
                {
                List<byte> data;
                List<uint> entryPointsInt;
                format.GetCodeAndEntryPoints(out data, out entryPointsInt);
                List<ulong> entryPoints = new List<ulong>();
                foreach (var e in entryPointsInt) // copy to ulong
                    entryPoints.Add(e);
                Disassemble(data, format.BitDepth, entryPoints, format.OptionalHeader.ImageBase);
                }
            } // DisassemblePE

        /// <summary>
        /// Show program usage
        /// </summary>
        static void Usage()
            {
            Console.WriteLine("{0} option", Environment.GetCommandLineArgs()[0]);
            Console.WriteLine("   where option is:");
            Console.WriteLine("   -t n    : test internal disassembler, ");
            Console.WriteLine("           : n=0 test all, n>0  test 1...n, n<0 test -n only.");
            Console.WriteLine("   -te n   : same as -t, show errors only");
            Console.WriteLine("   -p file : do a PE dump on the file");
            Console.WriteLine("   -d file : disassemble file, using main and exports as entry points");
            Console.WriteLine("   -s dir  : outfile: scan all files recursively from dir,");
            Console.WriteLine("           : disassemble all, gather stats, save in outfile");
            Console.WriteLine("Version {0}, Copyright Chris Lomont, July 2010, www.lomont.org", Version);
            }

        /// <summary>
        /// Do a PE file dump for the 
        /// </summary>
        static void DumpPE(List<byte> fileImage)
            {
            PEFormat format = new PEFormat();
            format.MessageEvent += (o, e) => Console.WriteLine(e.Message);
            bool success = format.Parse(fileImage);
            if (!success)
                Console.WriteLine("Failed");
            else
                {
                // Export table info
                Console.WriteLine("Exports table");
                var ex = format.ExportDirectory;
                Console.WriteLine(" Name            : ", ex.Name);
                Console.WriteLine(" Characteristics : {0:X8}", ex.Characteristics);
                Console.WriteLine(" TimeDateStamp   : {0:X8}", ex.TimeStamp);
                Console.WriteLine(" OrdinalBase     : {0:X8}", ex.Base);
                Console.WriteLine(" # of functions  : {0:X8}", ex.NumberOfFunctions);
                Console.WriteLine(" # of names      : {0:X8}", ex.NumberOfNames);
                Console.WriteLine();

                Console.WriteLine(" Entry Point   Ordinal   Name");
                foreach (var export in format.Exports)
                    {
                    Console.WriteLine("    {0:X8}      {1:X4}   {2}", export.AddressRVA, export.Ordinal, export.Name);
                    }


                Console.WriteLine("Success");
                }
            } // DumpPE

        /// <summary>
        /// Scan all valid PE files recursively from directory, gathering information.
        /// Output information to given filename.
        /// </summary>
        /// <param name="pathName"></param>
        static void ScanFiles(string pathName, string outfile)
            {

            ConsoleScanner scanner = new ConsoleScanner();
            scanner.Scan(pathName, outfile);
            }

        /// <summary>
        /// Read file if possible, else return null.
        /// </summary>
        /// <param name="filename"></param>
        /// <returns></returns>
        static List<byte> ReadFile(string filename)
            {
            if (!File.Exists(filename))
                {
                Console.WriteLine("Missing file {0}", filename);
                return null;
                }
            return File.ReadAllBytes(filename).ToList();
            }

        // todo - make tool: PE dump, ASM dump, and internal testing
        static void Main(string[] args)
            {

            if (args.Length < 1)
                {
                Usage();
                Environment.Exit(-1);
                }

            try
                {
                string command = args[0].ToLower();
                switch (command)
                    {
                    case "-d": // Disassemble
                        if (args.Length == 2)
                            {
                            Console.WriteLine("Disassembly of file {0}", args[1]);
                            DisassemblePE(ReadFile(args[1]));
                            }
                        else
                            {
                            Console.WriteLine("ERROR: Disassemble option needs only a filename after {0}", command);
                            Environment.Exit(-1);
                            }
                        break;
                    case "-p": // PE Dump
                        if (args.Length == 2)
                            {
                            Console.WriteLine("Dump of file {0}", args[1]);
                            DumpPE(ReadFile(args[1]));
                            }
                        else
                            {
                            Console.WriteLine("ERROR: Disassemble option needs only a filename after {0}", command);
                            Environment.Exit(-1);
                            }
                        break;
                    case "-s": // scanner
                        if (args.Length == 3)
                            {
                            Console.WriteLine("Scan of directory {0}", args[1]);
                            ScanFiles(args[1], args[2]);
                            }
                        else
                            {
                            Console.WriteLine("ERROR: Scan option needs a directory name and an output name after {0}", command);
                            Environment.Exit(-1);
                            }
                        break;
                    case "-t": // testing
                    case "-te": // testing, show only errors
                        // for testing option -t n where n is an integer describing the testing
                        // -te n shows error only
                        int testCount = -1;
                        if (Int32.TryParse(args[1], out testCount))
                            {
                            ConsoleDisassemblerTester tester = new ConsoleDisassemblerTester();
                            bool showErrorsOnly = args[0].ToLower() == "-te";
                            tester.Test(testCount, showErrorsOnly);
                            }
                        break;
                    default:
                        Usage();
                        Environment.Exit(-1);
                        break;
                    }
                }
            catch (Exception e)
                {
                ConsoleColor backupColor = Console.ForegroundColor;
                Console.ForegroundColor = ConsoleColor.Red;
                Console.WriteLine();
                Console.WriteLine("Fatal exception: {0} {1}",e.GetType() ,e.Message);
                Console.ForegroundColor = backupColor;
                }
            }
        }
    }
