using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.IO;
using System.Security.Cryptography;
using System.Runtime.InteropServices;
using System.Diagnostics;
using log4net.Config;


namespace HarvestDrivers
{
    class HarvestDrivers
    {
        private static readonly log4net.ILog log = log4net.LogManager.GetLogger(System.Reflection.MethodBase.GetCurrentMethod().DeclaringType);       
        static void copyDirectory(string strSource, string strDestination)
        {
            if (!Directory.Exists(strDestination))
            {
                Directory.CreateDirectory(strDestination);
            }
            DirectoryInfo dirInfo = new DirectoryInfo(strSource);
            FileInfo[] files = dirInfo.GetFiles();
            foreach (FileInfo tempfile in files)
            {
                if (tempfile.Extension !=".PNF" ) { 
                tempfile.CopyTo(Path.Combine(strDestination, tempfile.Name),true);
                }
            }
            DirectoryInfo[] dirctororys = dirInfo.GetDirectories();
            foreach (DirectoryInfo tempdir in dirctororys)
            {
                copyDirectory(Path.Combine(strSource, tempdir.Name), Path.Combine(strDestination, tempdir.Name));
            }

        }

        static void Main(string[] args)
        {
            XmlConfigurator.Configure(new System.IO.FileInfo(System.IO.Path.Combine(System.Environment.CurrentDirectory, "HarvestDrivers.exe.config")));
            SHA1 sha = new SHA1CryptoServiceProvider();

            List<string> loadeddrivers = new List<string>();
            ProcessStartInfo startInfo = new System.Diagnostics.ProcessStartInfo(@"LoadedDrivers.exe");
            startInfo.UseShellExecute = false;
            startInfo.RedirectStandardOutput = true;
            Process p = Process.Start(startInfo);
            string output = p.StandardOutput.ReadToEnd();
            p.WaitForExit();
            foreach (string drv in output.Split('\n'))
            {
                log.DebugFormat("Driver returned from psapi: {0}", drv);
                if (drv.Trim().Length > 0 )
                {
                    string filename;
                    if (drv.Substring(0, 4) == @"\??\")
                    {
                        filename = drv.Replace(@"\??\", @"");
                    }
                    else if (drv.Substring(0, 12) == @"\SystemRoot\")
                    {
                        filename = drv.Replace(@"\SystemRoot\",@"C:\WINDOWS\");
                    }
                    else
                    {
                        filename = "";
                    }

                    filename = filename.Replace("\r", "");
                    log.DebugFormat("Driver after processing: {0}", filename);
                    if (File.Exists(filename))
                    {
                        string hash = Convert.ToBase64String(sha.ComputeHash(System.IO.File.ReadAllBytes(filename)));
                        log.DebugFormat("File hash: {0}", hash);
                        loadeddrivers.Add(hash);
                    }
                    else
                    {
                        log.DebugFormat("File not found: {0}", filename);
                    }
                    
                }                
                                
            }

            List<string> oemfiles = new List<string>();
            string destinationDirectory="";
            if (args.Length == 0)
            {
                Console.Write("Destination Directory:");            
                destinationDirectory = Console.ReadLine();
            }
            else
            {
                destinationDirectory = args[0].TrimEnd('\\');
            }
            var oemFiles = Directory.EnumerateFiles("C:\\WINDOWS\\INF", "oem*.inf");
            foreach (string currentFile in oemFiles)
            {
                var fileinfo = new FileInfo(currentFile);
                log.DebugFormat("Found OEM*.INF: {0}", currentFile);
                string hash = Convert.ToBase64String(sha.ComputeHash(System.IO.File.ReadAllBytes(fileinfo.FullName)));
                log.DebugFormat("File hash: {0}", hash);
                oemfiles.Add(hash);                
            }
            Console.WriteLine(loadeddrivers.Count + " Drivers loaded in memory");
            Console.WriteLine(oemfiles.Count + " OEM Drivers Found");
            //now let's find all that drivers in FileRepository
            DirectoryInfo diTop = new DirectoryInfo(@"C:\Windows\System32\DriverStore\FileRepository");
            int cnt = 0;
            foreach (var di in diTop.EnumerateDirectories("*"))
            {
                bool found = false;
                string Driver="";
                string INF="";
                foreach (var fi in di.EnumerateFiles("*",SearchOption.AllDirectories))
                {
                    log.DebugFormat(" - {0}", fi.FullName);
                    if (loadeddrivers.Contains(Convert.ToBase64String(sha.ComputeHash(System.IO.File.ReadAllBytes(fi.FullName)))))                       
                    {
                        log.DebugFormat("{0} is a loaded driver", fi.FullName);
                        Driver = fi.Name;
                        loadeddrivers.Remove(Convert.ToBase64String(sha.ComputeHash(System.IO.File.ReadAllBytes(fi.FullName))));

                        foreach (var allfiles in di.EnumerateFiles("*.inf", SearchOption.AllDirectories))
                        {
                            log.DebugFormat(" - {0}", allfiles.FullName);
                            if (oemfiles.Contains(Convert.ToBase64String(sha.ComputeHash(System.IO.File.ReadAllBytes(allfiles.FullName)))))
                            {
                                log.DebugFormat("{0} is an OEM INF", allfiles.FullName);                                
                                INF = allfiles.Name;
                                oemfiles.Remove(Convert.ToBase64String(sha.ComputeHash(System.IO.File.ReadAllBytes(allfiles.FullName))));
                                found = true;
                                break;
                            }
                        }
                        break;
                    }
                    
                }
                if (found)
                {
                    Console.WriteLine("Copying: " + Path.Combine(di.Name) + " INF:" + INF + " Driver:" + Driver);
                    copyDirectory( Path.Combine(diTop.FullName, di.Name), Path.Combine(destinationDirectory, di.Name));
                    cnt++;
                }
            }

            Console.WriteLine(cnt + " Loaded OEM Drivers Copied");
            if (args.Length == 0)
            {
                Console.WriteLine("Press Enter to finish.");
                Console.ReadLine();
            }
        }

    }
}
