Merge pull request #1 from XKaguya/dev

Version 1.0.1

RDPInterceptor/API/Logger.cs

@@ -105,6 +105,8 @@ private static void LogAddLine(string message, SolidColorBrush color)
                 Paragraph paragraph = new Paragraph(new Run(message));
                 paragraph.Foreground = color;
                 logRichTextBox.Document.Blocks.Add(paragraph);
+
+                logRichTextBox.ScrollToEnd();
             });
         }
 

RDPInterceptor/API/NetworkInterceptor.cs

@@ -1,95 +1,112 @@
-using System.Collections.ObjectModel;
+using System;
+using System.Collections.Generic;
+using System.Collections.ObjectModel;
+using System.IO;
+using System.Linq;
 using System.Net;
+using System.Threading;
 using System.Threading.Tasks;
 using WindivertDotnet;
 
 namespace RDPInterceptor.API
 {
     public class NetworkInterceptor
     {
+        public static CancellationTokenSource CaptureCancellationTokenSource = new();
+
         public static bool IpWhitelistMode { get; set; } = true;
 
         public static ObservableCollection<IPAddress> IpAddrList { get; set; } = new();
 
-        public static bool IsCapture { get; set; } = true;
-
+        public static bool IsLogConnection { get; set; } = false;
+
+        public static bool WriteIntoLog { get; set; } = true;
+
+        public static ushort Port { get; set; } = 3389;
+
         private static WinDivert Divert { get; set; }
 
         private static WinDivertPacket Packet { get; set; }
-        
+
         private static WinDivertAddress Addr { get; set; }
 
-        public static bool AddIpIntoList(string Ip)
+        public static async Task<bool> AddIpIntoList(string Ip)
         {
             if (Ip == null)
             {
                 return false;
             }
 
-            IPAddress IpAddr = null;
+            IPAddress IpAddr;
 
             if (IPAddress.TryParse(Ip, out IpAddr))
             {
                 IpAddrList.Add(IpAddr);
+                await AddIpIntoWhitelistFile(IpAddr);
             }
             else
             {
                 Logger.Error($"Invalid IP/Domain.");
                 return false;
             }
-            
+
             return true;
         }
-        
-        public static async Task StartCapture()
+
+        public static async Task StartCapture(CancellationToken cancellationToken)
         {
             Logger.Log("Start Interceptor.");
-            
-            var filter = Filter.True.And(f => f.Tcp.DstPort == 3389);
-            
+
+            var filter = Filter.True.And(f => f.Tcp.DstPort == Port);
+
             Divert = new WinDivert(filter, WinDivertLayer.Network);
             Addr = new();
             Packet = new();
 
-            IsCapture = true;
-
-            while (IsCapture)
+            while (!cancellationToken.IsCancellationRequested)
             {
                 if (Divert != null)
                 {
                     try
                     {
-                        await Divert.RecvAsync(Packet, Addr);
+                        await Divert.RecvAsync(Packet, Addr, cancellationToken);
 
-                        if (ProcessPacket(Packet, Addr))
+                        if (await ProcessPacketAsync(Packet, Addr))
                         {
-                            await Divert.SendAsync(Packet, Addr);
+                            await Divert.SendAsync(Packet, Addr, cancellationToken);
                         }
                     }
-                    catch (TaskCanceledException ex)
+                    catch (OperationCanceledException ex)
                     {
                         Logger.Log($"Stop Interceptor.");
-                        IsCapture = false;
                     }
                 }
             }
         }
 
-        public static async Task StopCapture()
+        public static async Task StopCapture(CancellationToken cancellationToken)
         {
-            IsCapture = false;
-
-            if (Divert != null)
+            if (cancellationToken.CanBeCanceled)
             {
-                Divert.Dispose();
-                Packet.Dispose();
-                Addr.Dispose();
+                CaptureCancellationTokenSource?.Cancel();
             }
-
+            else
+            {
+                Logger.Error("You cannot stop capture at this time.");
+            }
+
             Logger.Log("Capture has now stopped.");
         }
 
-        private unsafe static bool ProcessPacket(WinDivertPacket Packet, WinDivertAddress Address)
+        private static unsafe void GetIpAddresses(IPV4Header* header, out IPAddress srcIpAddr, out IPAddress dstIpAddr)
+        {
+            IPAddress srcIp = header->SrcAddr;
+            IPAddress dstIp = header->DstAddr;
+            srcIpAddr = srcIp;
+            dstIpAddr = dstIp;
+        }
+
+        public static async Task<bool> ProcessPacketAsync(WinDivertPacket Packet, WinDivertAddress Address)
         {
             if (Packet != null)
             {
@@ -100,31 +117,159 @@ private unsafe static bool ProcessPacket(WinDivertPacket Packet, WinDivertAddres
                 else
                 {
                     var result = Packet.GetParseResult();
-                    IPAddress SrcIpAddr = result.IPV4Header->SrcAddr;
-                    IPAddress DstIpAddr = result.IPV4Header->DstAddr;
-
+                    IPAddress SrcIpAddr, DstIpAddr;
+                    unsafe
+                    {
+                        GetIpAddresses(result.IPV4Header, out SrcIpAddr, out DstIpAddr);
+                    }
+
                     if (IpAddrList.Contains(SrcIpAddr))
                     {
-                        Logger.Debug($"Incomming RDP Connection from {SrcIpAddr} has been accepted.");
+                        LogConnections(IsLogConnection,$"Incoming RDP Connection from {SrcIpAddr} has been accepted.");
                         Packet.CalcChecksums(Address);
                         return true;
                     }
                     else if (IpAddrList.Contains(DstIpAddr))
                     {
-                        Logger.Debug($"Outgoing RDP Connection to {DstIpAddr} has been accepted.");
+                        LogConnections(IsLogConnection,$"Outgoing RDP Connection to {DstIpAddr} has been accepted.");
                         Packet.CalcChecksums(Address);
                         return true;
                     }
                     else
                     {
-                        Logger.Debug($"Incomming RDP Connection from {SrcIpAddr} has been refused.");
+                        LogConnections(IsLogConnection, $"Incoming RDP Connection from {SrcIpAddr} has been refused.");
+                        await LogConnectionAsync(SrcIpAddr);
+
                         return false;
                     }
                 }
             }
 
             return false;
         }
+
+        private static void LogConnections(bool isLogConnection, string content)
+        {
+            if (isLogConnection)
+            {
+                Logger.Debug(content);
+            }
+        }
+
+        private static async Task LogConnectionAsync(IPAddress srcIpAddr)
+        {
+            string logFilePath = "Connectionlist.log";
+
+            if (File.Exists(logFilePath))
+            {
+                string[] lines = await File.ReadAllLinesAsync(logFilePath);
+                if (lines.Contains(srcIpAddr.ToString()))
+                {
+                    return;
+                }
+            }
+            else
+            {
+                FileStream fs = File.Create(logFilePath);
+                fs.Close();
+            }
+
+            using (StreamWriter writer = File.AppendText(logFilePath))
+            {
+                await writer.WriteLineAsync(srcIpAddr.ToString());
+            }
+        }
+
+        private static readonly SemaphoreSlim semaphore = new(1);
+
+        public static async void ReadLinesFromFileAsync()
+        {
+            await semaphore.WaitAsync();
+
+            string WhitelistFilePath = "Whitelist.txt";
+
+            try
+            {
+                if (File.Exists(WhitelistFilePath))
+                {
+                    IPAddress IpAddr;
+
+                    foreach (string ip in await File.ReadAllLinesAsync(WhitelistFilePath))
+                    {
+                        if (IPAddress.TryParse(ip, out IpAddr))
+                        {
+                            IpAddrList.Add(IpAddr);
+
+                            Logger.Log($"IP {ip} has been read into whitelist.");
+                        }
+                        else
+                        {
+                            Logger.Error($"ERROR! Failed to parse {ip}.");
+                        }
+                    }
+                }
+                else
+                {
+                    FileStream fs = File.Create(WhitelistFilePath);
+                    fs.Close();
+                }
+            }
+            finally
+            {
+                semaphore.Release();
+            }
+        }
+
+        public static async Task AddIpIntoWhitelistFile(IPAddress ipAddress)
+        {
+            Logger.Debug($"Method AddIpIntoWhitelistFile called.");
+
+            try
+            {
+                string WhitelistFilePath = "Whitelist.txt";
+
+                if (File.Exists(WhitelistFilePath))
+                {
+                    Logger.Debug($"File {WhitelistFilePath} exists. Proceeding...");
+
+                    await semaphore.WaitAsync();
+
+                    string[] lines = await File.ReadAllLinesAsync(WhitelistFilePath);
+                    if (Array.Exists(lines, line => line.Equals(ipAddress.ToString())))
+                    {
+                        Logger.Debug($"IP {ipAddress} already in {WhitelistFilePath}");
+                        return;
+                    }
+                    else
+                    {
+                        using (StreamWriter writer = File.AppendText(WhitelistFilePath))
+                        {
+                            await writer.WriteLineAsync(ipAddress.ToString());
+                        }
+
+                        Logger.Debug($"IP {ipAddress} has been written into {WhitelistFilePath}");
+                    }
+                }
+                else
+                {
+                    FileStream fs = File.Create(WhitelistFilePath);
+                    fs.Close();
+
+                    Logger.Error($"File {WhitelistFilePath} doesn't exist. Now create file {WhitelistFilePath}.");
+
+                    await AddIpIntoWhitelistFile(ipAddress);
+                }
+            }
+            catch (Exception e)
+            {
+                Logger.Error(e.Message + e.Message);
+                throw;
+            }
+            finally
+            {
+                semaphore.Release();
+            }
+        }
     }
 }