using System;
using System.Diagnostics;
using System.Collections.Generic;
using System.Security.Cryptography;
using System.IO;
using System.Text;
using NDesk.Options;

/*  Description:
        Light-weight utility to help maintain an archive of Windows system files.
 *      Possibly useful for binary diff analysis of MS updates.
 *      
 *  Distributed:
 *      L1pht Offensive Labs            http://www.l1pht.com
 *      Bugs, comments, suggestions     saintpatrick@l1pht.com
 *      
 *  Feature Requests:
 *      - Store MD5 information for each file in a searchable database. Neat, but not certain the need is there.
 *      
 *  History:
 *      - 3/20/2010     Built
 *	- 3/21/2010		Fixed a bug so that it's possible to do crazy stuff like have spaces in folder names ;P (QA win)
*/

namespace binaryeti
{
    class Program
    {
        // This should work across versions
        static string systemFolder = Environment.GetFolderPath(Environment.SpecialFolder.System);

        static void Main(string[] args)
        {
            string homepath=string.Empty;
            bool help = false;

            OptionSet yetiOptions = new OptionSet()
            {
                {"p|path=","The parent path of the archive",v=>homepath=v},
                {"h|help","Show this screen",v => help=v != null},
            };

            // Let the Yeti have some args. nom nom.
            yetiOptions.Parse(args);

            if (help)
            {
                ShowHelp(yetiOptions);
                return;
            };

            // If we have a parent we can do the rest.
            if (homepath == string.Empty)
            {
                ShowHelp(yetiOptions);
            }
            else
            {
                

                DirectoryInfo systemDir = new DirectoryInfo(systemFolder);
                string sanitizedFileVersion;

                Console.WriteLine("[*]  Starting archive process");

                foreach (FileInfo file in systemDir.GetFiles())
                {
                    if (file.Extension == ".dll" | file.Extension == ".ocx")
                    {
                        string newParentPath = homepath + "\\" + file.Name;
                        newParentPath = System.Text.RegularExpressions.Regex.Replace(newParentPath, @"""", "");

                        FileVersionInfo currentFilesVersionInfo = FileVersionInfo.GetVersionInfo(file.FullName);
                        if (currentFilesVersionInfo.FileVersion != null)
                        {
                            sanitizedFileVersion = FilenameStripperPole(currentFilesVersionInfo.FileVersion);
                        }
                        else
                        {
                            // If I don't have some sort of reliable version I'm using the MD5 of the file.
                            sanitizedFileVersion = GetMD5HashFromFile(file.FullName);
                            Console.WriteLine("[*]  Had to use MD5 for version folder for file: " + file.Name);
                        }
                       
                        if (Directory.Exists(newParentPath)!=true)
                        {
                            Directory.CreateDirectory(newParentPath);
                        }
                        try
                        {
                            if (Directory.Exists(newParentPath + "\\" + sanitizedFileVersion) != true)
                            {
                                Directory.CreateDirectory(newParentPath + "\\" + sanitizedFileVersion);
                                Console.WriteLine("[*]  Archiving " + newParentPath + "\\" + sanitizedFileVersion);
                                File.Copy(file.FullName, newParentPath + "\\" + sanitizedFileVersion + "\\" + file.Name);
                            }
                        }
                        catch (Exception ex)
                        {
                            Console.WriteLine("[X] Error archiving due to: " + ex.Message);
                        }

                    }
                }

                Console.WriteLine("[*]  Archive process finished");

            }
        }

        
        /// <summary>
        /// Shows the command line options for the utility
        /// </summary>
        /// <param name="myoptionset">The NDesk OptionSet created for the utility</param>
        static void ShowHelp (OptionSet myoptionset)
        {
            Console.WriteLine ("Usage: binaryeti [OPTIONS]");
            Console.WriteLine();
            Console.WriteLine ("Fill out the args and I'll archive your "+ systemFolder +" dlls and ocxs.");
            Console.WriteLine ();
            Console.WriteLine ("Options:");
            myoptionset.WriteOptionDescriptions (Console.Out);
        }

        /// <summary>
        /// Function to retrieve MD5 hash for a file
        /// </summary>
        /// <param name="fullFileName">The path of the file to hash</param>
        /// <returns>A string containing </returns>
        static string GetMD5HashFromFile(string fullFileName)
        {
            StringBuilder sb = new StringBuilder();

            try
            {
                FileStream file = new FileStream(fullFileName, FileMode.Open);
                MD5 md5 = new MD5CryptoServiceProvider();
                byte[] retVal = md5.ComputeHash(file);
                file.Close();

                for (int i = 0; i < retVal.Length; i++)
                {
                    sb.Append(retVal[i].ToString("X2"));
                }
            }
            catch (Exception ex)
            {
                Console.WriteLine("[X] Error during MD5 hashing of: " + fullFileName);
                Console.WriteLine("[X] " + ex.Message);
            }
            return sb.ToString();
        }

        /// <summary>
        /// Lame function to strip out invalid directory chars.
        /// </summary>
        /// <param name="originalString"></param>
        /// <returns></returns>
        static string FilenameStripperPole(string originalString)
        {
            originalString = originalString.Replace(":", "");
            originalString = originalString.Replace("?", "");
            originalString = originalString.Replace("<", "");
            originalString = originalString.Replace(">", "");
            originalString = originalString.Replace("|", "");
            return originalString;
        }
    }
    
}