CustomNetworkProcessMonitoring/ETWwrapper.cs at master · VodOleg/CustomNetworkProcessMonitoring · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
using Microsoft.Diagnostics.Tracing.Parsers;
using Microsoft.Diagnostics.Tracing.Session;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

namespace ProcessesMonitoring
{

    public static class ETWwrapper
    {
        private static Task m_workingThread;
        private static CancellationTokenSource _tokenSource = null;
        private static Dictionary<int, SessionData> m_ProcessIDs = new Dictionary<int, SessionData>();
        private static TraceEventSession m_EtwSession;

        public static void start()
        {
            _tokenSource = new CancellationTokenSource();
            var cancelToken = _tokenSource.Token;
            m_workingThread = Task.Run(() => StartEtwSession(cancelToken));
        }

        public static void addProcess(int processID)
        {
            lock (m_ProcessIDs)
            {
                m_ProcessIDs[processID] = new SessionData();
            }
        }

        private static void StartEtwSession(CancellationToken cancelToken)
        {
            try
            {

                using (m_EtwSession = new TraceEventSession("MyKernelAndClrEventsSession"))
                {
                    m_EtwSession.EnableKernelProvider(KernelTraceEventParser.Keywords.NetworkTCPIP);

                    m_EtwSession.Source.Kernel.TcpIpRecv += data =>
                    {
                        //Console.WriteLine($"process reading: {data.ProcessID}");
                        if (cancelToken.IsCancellationRequested)
                        {
                            // stop the monitor
                            Console.WriteLine("canceled rec");
                            return;
                        }
                        if (m_ProcessIDs.ContainsKey(data.ProcessID))
                        {
                            //Console.WriteLine($"rec: {data.ProcessID} ({data.ProcessName}) ->{data.size}");
                            lock (m_ProcessIDs[data.ProcessID])
                            {
                                m_ProcessIDs[data.ProcessID].received = data.size;
                                m_ProcessIDs[data.ProcessID].timestamp = data.TimeStamp;
                            }

                        }
                        else
                        {

                        }
                    };

                    m_EtwSession.Source.Kernel.TcpIpDisconnect += data =>
                    {
                        if (m_ProcessIDs.ContainsKey(data.ProcessID))
                        {
                            Console.WriteLine($"detected {data.ProcessID} tcp disconnect");
                            lock (m_ProcessIDs[data.ProcessID])
                            {
                                m_ProcessIDs[data.ProcessID].disconnected = true;
                                m_ProcessIDs[data.ProcessID].received = data.size;
                                m_ProcessIDs[data.ProcessID].timestamp = data.TimeStamp;
                            }
                        }
                    };

                    m_EtwSession.Source.Kernel.TcpIpSend += data =>
                    {
                        if (cancelToken.IsCancellationRequested)
                        {
                            // stop the monitor
                            return;
                        }
                        if (m_ProcessIDs.ContainsKey(data.ProcessID))
                        {
                            lock (m_ProcessIDs[data.ProcessID])
                            {
                                m_ProcessIDs[data.ProcessID].sent = data.size;
                            }
                        }
                    };

                    m_EtwSession.Source.Process();
                }
            }
            catch
            {
                //ResetCounters(); // Stop reporting figures
                // Probably should log the exception
                Console.WriteLine("Exception caught in session ETW");
            }
        }

        public static SessionData getData(int ProcessID)
        {
            SessionData dataToReturn = new SessionData();


            if (m_ProcessIDs.ContainsKey(ProcessID))
            {
                lock (m_ProcessIDs[ProcessID])
                {
                    dataToReturn.hardCopy(m_ProcessIDs[ProcessID]);
                }
            }
            return dataToReturn;
        }
    }
}