CollabVM-Agent-Windows/CollabVMAgent/Program.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.IO;
using System.Threading;

using CollabVMAgent.Protocol;
using MsgPack.Serialization;
using System.Security.Principal;
using System.Diagnostics;

namespace CollabVMAgent
{
    class Program
    {

        static readonly string UploadsFolder = Environment.GetEnvironmentVariable("USERPROFILE") + @"\Desktop";
        static MessagePackSerializer serializer = MessagePackSerializer.Get<ProtocolMessage>();
        static VirtIOSerial vio;
        static void Main(string[] args)
        {
            if (!IsAdmin())
            {
                Process cur = Process.GetCurrentProcess();
                Process prc = new Process();
                prc.StartInfo.Verb = "runas";
                prc.StartInfo.FileName = cur.MainModule.FileName;
                prc.Start();
                Environment.Exit(0);
            }
#if DEBUG
            Console.WriteLine("Starting CollabVM Agent in Debug mode");
#endif
            vio = new VirtIOSerial();
            vio.Data += OnData;
            if (vio.WriteMsg(new ProtocolMessage{
                Operation = ProtocolOperation.NOP
            }))
            {
#if DEBUG
                Console.WriteLine("Wrote NOP");
#endif
            }
            Thread.Sleep(-1);
        }

        private static void OnData(object sender, TypedEventArgs<byte[]> e)
        {
            ProtocolMessage msg;
            using (MemoryStream ms = new MemoryStream(e.Data))
            {
                msg = (ProtocolMessage)serializer.Unpack(ms);
            }
            switch (msg.Operation)
            {
                case ProtocolOperation.UploadFile:
                    {
                        foreach (char c in Path.GetInvalidFileNameChars())
                        {
                            if (msg.Filename.Contains(c))
                            {
                                Console.Error.WriteLine($"Invalid filename \"{msg.Filename}\"");
                                return;
                            }
                        }
                        try
                        {
                            File.WriteAllBytes(UploadsFolder + @"\" + msg.Filename, msg.FileData);
                        } catch (Exception ex)
                        {
#if DEBUG
                            Console.WriteLine($"Failed to save {msg.Filename}: {ex.Message}");
#endif
                            return;
                        }
#if DEBUG
                        Console.WriteLine($"Successfully wrote {msg.Filename} ({msg.FileData.Length} bytes)");
#endif
                        new Thread(() =>
                        {
                            try
                            {
                                Process.Start(UploadsFolder + @"\" + msg.Filename).Start();
                            } catch (Exception ex)
                            {
#if DEBUG
                                Console.WriteLine($"Error while starting {msg.Filename}: {ex.Message}");
#endif
                            }
                        }).Start();
                        break;
                    }
                case ProtocolOperation.NOP:
                    {
#if DEBUG
                        Console.WriteLine("Got NOP from server");
#endif
                        if (vio.WriteMsg(new ProtocolMessage
                        {
                            Operation = ProtocolOperation.NOP
                        }))
                        {
#if DEBUG
                            Console.WriteLine("Wrote NOP");
#endif
                        }
                        break;
                    }

            }
        }

        static bool IsAdmin()
        {
            WindowsIdentity identity = WindowsIdentity.GetCurrent();
            WindowsPrincipal principal = new WindowsPrincipal(identity);
            bool isElevated = principal.IsInRole(WindowsBuiltInRole.Administrator);
            identity.Dispose();
            return isElevated;
        }
    }
}