SimpleVirtualDrive/SvdCli/CliProgram.cs

using SvdCli.App;
using SvdCli.ServiceUtils;
using SvdCli.Storage;
using SvdCli.WindowsService;
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Management;
using System.Net;
using System.Reflection;
using System.Security;
using System.ServiceProcess;
using System.Text.RegularExpressions;
using SvdCli.DiskUtils;
using SvdCli.Storage.Implements;

namespace SvdCli
{
    internal static class CliProgram
    {
        private static void Main(string[] args)
        {
            if (args.Length < 2)
            {
                Console.WriteLine(@"Usage:");
                Console.WriteLine(@"");
                Console.WriteLine(@"Create Raw 50GB C:\Path\To\Your\Image.img");
                Console.WriteLine(@"Create Bdd 10GB 64KB C:\Path\To\Your\Bdd.bdd");
                Console.WriteLine(@"Create Bdd 64KB C:\Path\To\Your\Image.img C:\Path\To\Your\Bdd.bdd");
                Console.WriteLine(@"Compact Bdd C:\Path\To\Your\Bdd.bdd");
                Console.WriteLine(@"");
                Console.WriteLine(@"[Install ServiceName|Service] <Mount|Server 0.0.0.0 2333> RamDisk 10GB");
                Console.WriteLine(@"[Install ServiceName|Service] <Mount|Server 0.0.0.0 2333> NtfsRamDisk 10GB");
                Console.WriteLine(@"[Install ServiceName|Service] <Mount|Server 0.0.0.0 2333> TempNtfsRamDisk 10GB");
                Console.WriteLine(@"[Install ServiceName|Service] <Mount|Server 0.0.0.0 2333> RamDisk 10GB C:\Path\To\Your\Bdd.Bdd [ReadOnly]");
                Console.WriteLine(@"[Install ServiceName|Service] <Mount|Server 0.0.0.0 2333> Image C:\Path\To\Your\Image.img [ReadOnly]");
                Console.WriteLine(@"[Install ServiceName|Service] <Mount|Server 0.0.0.0 2333> Bdd C:\Path\To\Your\Bdd.bdd [ReadOnly]");
                Console.WriteLine(@"[Install ServiceName|Service] Mount Net 192.168.233.233 2333");
                Console.WriteLine(@"[Install ServiceName|Service] Mount ZTempNtfsRamDisk 10GB");
                Console.WriteLine(@"[Install ServiceName|Service] Server 192.168.233.233 2333 Dispatch 64KB C:\Path\To\Your\Image.img C:\Path\To\Your\Dir\");
                Console.WriteLine(@"");
                Console.WriteLine(@"Uninstall ServiceName           Safe:Restrict to register by this program");
                Console.WriteLine(@"");
                Console.Write("Press ENTER to exit...");
                Console.ReadLine();
                return;
            }

            EndPoint serverEndPoint = null;
            ISvdStorage storage = null;
            switch (args[0].ToLower())
            {
                case "install":
                    {
                        var serviceName = args[1];
                        var serviceArgs = new List<string>() { Assembly.GetExecutingAssembly().Location, "Service" };
                        serviceArgs.AddRange(args.Skip(2));
                        ServiceHelper.Install(serviceName, "", string.Join(" ", serviceArgs), "", ServiceStartType.Auto);
                    }
                    break;

                case "uninstall":
                    {
                        var serviceName = args[1];
                        var binPath = ServiceHelper.GetBinPath(serviceName);

                        if (true == binPath?.Contains(Path.GetFileName(Assembly.GetExecutingAssembly().Location))) ServiceHelper.Uninstall(serviceName);
                        else throw new SecurityException("Specified service must register by this program");
                    }
                    break;

                case "service":
                    args = args.Skip(1).ToArray();

                    switch (args[0].ToLower())
                    {
                        default: throw new InvalidOperationException("Invalid param: OP");
                        case "server":
                            serverEndPoint = new IPEndPoint(IPAddress.Parse(args[1]), int.Parse(args[2]));
                            if (args[3].ToLower() == "dispatch")
                            {
                                //Start TCP Server and listening
                                //Accept and Get Client MAC
                                //Create Bdd If Not Exist
                                //Serving Bdd in `forked' Process use Socket.DuplicateAndClose
                                throw new NotImplementedException();
                            }
                            args = args.Skip(3).ToArray();
                            goto case "mount";
                        case "mount":
                            storage = ParseMount(args.Skip(1).ToArray());
                            break;
                    }

                    ServiceBase.Run(new ServiceBase[] { new SvdServiceForWindows(new SvdService(storage, serverEndPoint)) });
                    break;

                default:
                    switch (args[0].ToLower())
                    {
                        case "create":
                            ParseCreate(args.Skip(1).ToArray());
                            return;

                        case "compact":
                            ParseCompact(args.Skip(1).ToArray());
                            return;

                        case "server":
                            serverEndPoint = new IPEndPoint(IPAddress.Parse(args[1]), int.Parse(args[2]));
                            args = args.Skip(3).ToArray();

                            goto case "mount";
                        case "mount":
                            storage = ParseMount(args.Skip(1).ToArray());
                            break;
                    }

                    var svc = new SvdService(storage, serverEndPoint);
                    svc.Start();
                    Console.WriteLine();
                    Console.Write("Press ENTER to Stop...");
                    Console.ReadLine();
                    Console.Write("Stopping...");
                    svc.Stop();
                    Console.Write("Stopped.");
                    break;
            }
        }

        private static void ParseCreate(string[] args)
        {
            // 0   1    2    3
            // Bdd 10GB 64KB C:\Path\To\Your\Bdd.bdd
            // Bdd 64KB C:\Path\To\Your\Image.img C:\Path\To\Your\Bdd.bdd
            // Raw 50GB C:\Path\To\Your\Image.img

            if (false == ParseDiskSize(args[1], out var value1, out var unit1)) throw new ArgumentException("Failure to parse disk size");
            var size1 = CalcDiskSize(value1, unit1);

            switch (args[0].ToLower())
            {
                case "raw":
                    if (File.Exists(args[2])) throw new IOException("File Already Exist");
                    {
                        using var img = File.Create(args[2]);
                        img.SetSparse();
                        img.SetLength(size1);
                    }
                    break;

                case "bdd":
                    if (args.Length > 3)
                    {
                        if (ParseDiskSize(args[2], out var value2, out var unit2)
                            && ".bdd" == Path.GetFileName(args[2])?.ToLower()
                            && false == File.Exists(args[3]))
                        {
                            //bdd 2=size 3=bdd[not exist]

                            var size = size1;

                            var blockSize = CalcDiskSize(value2, unit2);

                            throw new NotImplementedException();
                        }

                        if (".img" == Path.GetExtension(args[2])?.ToLower()
                            && File.Exists(args[2])
                            && ".bdd" == Path.GetExtension(args[3])?.ToLower()
                            && false == File.Exists(args[3]))
                        {
                            //snapshot 2=img[exist] 3=bdd[not exist]

                            var blockSize = value1;

                            throw new NotImplementedException();
                        }
                    }

                    throw new ArgumentException("Unknown bdd create action");

                default: throw new ArgumentException("Unknown image type");
            }
        }

        private static ISvdStorage ParseMount(string[] args)
        {
            // 0       1    2                       3
            // RamDisk 10GB C:\Path\To\Your\Bdd.Bdd [ReadOnly]
            // RamDisk 10GB
            // NtfsRamDisk 10GB
            // TempNtfsRamDisk 10GB
            // Image C:\Path\To\Your\Image.img [ReadOnly]
            // Bdd C:\Path\To\Your\Bdd.bdd [ReadOnly]
            // Net 192.168.233.233 2333

            var op = args[0].ToLower();
            switch (op)
            {
                case "ramdisk":
                case "ntfsramdisk":
                case "tempntfsramdisk":
                case "ztempntfsramdisk":
                    if (ParseDiskSize(args[1], out var value, out var unit))
                    {
                        if (false == string.IsNullOrWhiteSpace(unit) && false == unit.StartsWith("g"))
                            throw new ArgumentException("RamDisk only GB supported");
                        var storage = new UnmanagedGigabyteRamDisk(value);

                        //TODO: load bdd
                        if (args.Length > 2) throw new NotImplementedException();

                        if (args.Length > 3 && args[3].ToLower() == "readonly") storage.WriteProtect = true;


                        if (op == "ntfsramdisk") FsMaker.MakeNtfs(storage, "RamDisk");
                        if (op == "tempntfsramdisk") FsMaker.MakeNtfs(storage, "RamDisk", n => n.CreateDirectory("Temp"));
                        if (op == "ztempntfsramdisk")
                        {
                            var guid = Guid.NewGuid().ToString("N");
                            FsMaker.MakeNtfs(storage, guid, n => n.CreateDirectory("Temp"));
                            storage.Mounted += delegate
                            {
                                var volumes = new ManagementObjectSearcher("root\\CIMV2", "SELECT * FROM Win32_Volume");
                                foreach (ManagementObject volume in volumes.Get())
                                {
                                    volume.Get();
                                    if (volume["Label"] is string str && str == guid)
                                    {
                                        volume["Label"] = "RamDisk";
                                        volume["DriveLetter"] = "Z:";
                                        volume.Put();
                                        break;
                                    }
                                }
                            };
                        }

                        return storage;
                    }

                    throw new ArgumentException("Failure to parse args");

                case "image":
                    {
                        var storage = new RawImageStorage(args[1], args.Length > 2 && "readonly" == args[2].ToLower());
                        return storage;
                    }
                case "bdd":
                case "net":
                default:
                    throw new ArgumentOutOfRangeException();
            }
        }

        private static void ParseCompact(string[] args)
        {
            //Bdd C:\Path\To\Your\Bdd.bdd
            throw new NotImplementedException();
        }

        //////////////////////////////////////////////

        private static readonly Regex DiskSizeRegex = new Regex(@"^(\d+)([KkMmGg]{1}[Bb]{0,1}){0,1}$", RegexOptions.Compiled);

        private static bool ParseDiskSize(string input, out int value, out string unit)
        {
            var m = DiskSizeRegex.Match(input);
            if (false == m.Success)
            {
                value = 0;
                unit = null;
                return false;
            }

            unit = m.Groups[2].Value.ToLower();
            value = int.Parse(m.Groups[1].Value);
            return true;
        }

        private static long CalcDiskSize(int value, string unit)
        {
            switch (unit.ToLower())
            {
                //case "":
                //    return value;

                case "k":
                case "kb":
                    return value * CapacityUnits.KiloByte;

                case "m":
                case "mb":
                    return value * CapacityUnits.MegaByte;

                case "g":
                case "gb":
                    return value * CapacityUnits.Gigabyte;

                default: throw new ArgumentOutOfRangeException();
            }
        }
    }
}