Private/Tcp445.ps1
|
#Requires -Version 5.1 function Initialize-SMBeatTcp445Native { if ('SMBeat.Native.Tcp445' -as [type]) { return } $code = @' using System; using System.Collections.Generic; using System.Net; using System.Runtime.InteropServices; namespace SMBeat.Native { public sealed class Tcp445Conn { public int Af; public string LocalIp; public int LocalPort; public string RemoteIp; public int RemotePort; public uint State; public uint LocalAddrV4; public uint RemoteAddrV4; public byte[] LocalAddrV6; public byte[] RemoteAddrV6; public uint LocalScopeId; public uint RemoteScopeId; public ulong BytesOut; public ulong BytesIn; public bool ReadOk; } public static class Tcp445 { private const int AF_INET = 2; private const int AF_INET6 = 23; private const int TCP_TABLE_OWNER_PID_CONNECTIONS = 4; private const uint ERROR_INSUFFICIENT_BUFFER = 122; private const int TcpConnectionEstatsData = 1; private const int MIB_TCP_STATE_ESTAB = 5; private const int SMB_PORT = 445; [DllImport("iphlpapi.dll", SetLastError = true)] private static extern uint GetExtendedTcpTable( IntPtr pTcpTable, ref int pdwSize, bool bOrder, int ulAf, int TableClass, int Reserved); [DllImport("iphlpapi.dll", SetLastError = true)] private static extern uint SetPerTcpConnectionEStats( IntPtr Row, int EstatsType, byte[] Rw, uint RwVersion, uint RwSize, uint Offset); [DllImport("iphlpapi.dll", SetLastError = true)] private static extern uint GetPerTcpConnectionEStats( IntPtr Row, int EstatsType, IntPtr Rw, uint RwVersion, uint RwSize, IntPtr Ros, uint RosVersion, uint RosSize, byte[] Rod, uint RodVersion, uint RodSize); [DllImport("iphlpapi.dll", SetLastError = true)] private static extern uint SetPerTcp6ConnectionEStats( IntPtr Row, int EstatsType, byte[] Rw, uint RwVersion, uint RwSize, uint Offset); [DllImport("iphlpapi.dll", SetLastError = true)] private static extern uint GetPerTcp6ConnectionEStats( IntPtr Row, int EstatsType, IntPtr Rw, uint RwVersion, uint RwSize, IntPtr Ros, uint RosVersion, uint RosSize, byte[] Rod, uint RodVersion, uint RodSize); private static int HostPort(uint netPort) { int p = (int)(netPort & 0xFFFF); return ((p >> 8) & 0xFF) | ((p & 0xFF) << 8); } private static bool IsLoopbackV4(uint addr) { return (addr & 0xFF) == 127; } private static bool IsLoopbackV6(byte[] addr) { if (addr == null || addr.Length < 16) { return false; } for (int i = 0; i < 15; i++) { if (addr[i] != 0) { return false; } } return addr[15] == 1; } public static Tcp445Conn[] ListEstablished() { List<Tcp445Conn> list = new List<Tcp445Conn>(); AddV4(list); AddV6(list); return list.ToArray(); } public static bool Enable(Tcp445Conn conn) { if (conn == null) { return false; } byte[] rw = new byte[] { 1 }; IntPtr row = AllocRow(conn); try { uint rc; if (conn.Af == AF_INET6) { rc = SetPerTcp6ConnectionEStats(row, TcpConnectionEstatsData, rw, 0, 1, 0); } else { rc = SetPerTcpConnectionEStats(row, TcpConnectionEstatsData, rw, 0, 1, 0); } return rc == 0; } finally { Marshal.FreeHGlobal(row); } } public static bool TryRead(Tcp445Conn conn) { if (conn == null) { return false; } byte[] rod = new byte[96]; IntPtr row = AllocRow(conn); try { uint rc; if (conn.Af == AF_INET6) { rc = GetPerTcp6ConnectionEStats(row, TcpConnectionEstatsData, IntPtr.Zero, 0, 0, IntPtr.Zero, 0, 0, rod, 0, 96); } else { rc = GetPerTcpConnectionEStats(row, TcpConnectionEstatsData, IntPtr.Zero, 0, 0, IntPtr.Zero, 0, 0, rod, 0, 96); } if (rc != 0) { conn.ReadOk = false; return false; } conn.BytesOut = BitConverter.ToUInt64(rod, 0); conn.BytesIn = BitConverter.ToUInt64(rod, 8); conn.ReadOk = true; return true; } finally { Marshal.FreeHGlobal(row); } } private static void AddV4(List<Tcp445Conn> list) { int size = 0; uint rc = GetExtendedTcpTable(IntPtr.Zero, ref size, false, AF_INET, TCP_TABLE_OWNER_PID_CONNECTIONS, 0); if (rc != ERROR_INSUFFICIENT_BUFFER || size <= 0) { return; } IntPtr buf = Marshal.AllocHGlobal(size); try { rc = GetExtendedTcpTable(buf, ref size, false, AF_INET, TCP_TABLE_OWNER_PID_CONNECTIONS, 0); if (rc != 0) { return; } int count = Marshal.ReadInt32(buf); int offset = 4; for (int i = 0; i < count; i++) { uint state = (uint)Marshal.ReadInt32(buf, offset); uint localAddr = (uint)Marshal.ReadInt32(buf, offset + 4); uint localPort = (uint)Marshal.ReadInt32(buf, offset + 8); uint remoteAddr = (uint)Marshal.ReadInt32(buf, offset + 12); uint remotePort = (uint)Marshal.ReadInt32(buf, offset + 16); offset += 24; if ((int)state != MIB_TCP_STATE_ESTAB) { continue; } if (HostPort(localPort) != SMB_PORT) { continue; } if (IsLoopbackV4(localAddr) || IsLoopbackV4(remoteAddr)) { continue; } Tcp445Conn c = new Tcp445Conn(); c.Af = AF_INET; c.State = state; c.LocalAddrV4 = localAddr; c.RemoteAddrV4 = remoteAddr; c.LocalPort = HostPort(localPort); c.RemotePort = HostPort(remotePort); c.LocalIp = new IPAddress((long)localAddr).ToString(); c.RemoteIp = new IPAddress((long)remoteAddr).ToString(); list.Add(c); } } finally { Marshal.FreeHGlobal(buf); } } private static void AddV6(List<Tcp445Conn> list) { int size = 0; uint rc = GetExtendedTcpTable(IntPtr.Zero, ref size, false, AF_INET6, TCP_TABLE_OWNER_PID_CONNECTIONS, 0); if (rc != ERROR_INSUFFICIENT_BUFFER || size <= 0) { return; } IntPtr buf = Marshal.AllocHGlobal(size); try { rc = GetExtendedTcpTable(buf, ref size, false, AF_INET6, TCP_TABLE_OWNER_PID_CONNECTIONS, 0); if (rc != 0) { return; } int count = Marshal.ReadInt32(buf); int offset = 4; for (int i = 0; i < count; i++) { byte[] localAddr = new byte[16]; Marshal.Copy(IntPtr.Add(buf, offset), localAddr, 0, 16); uint localScope = (uint)Marshal.ReadInt32(buf, offset + 16); uint localPort = (uint)Marshal.ReadInt32(buf, offset + 20); byte[] remoteAddr = new byte[16]; Marshal.Copy(IntPtr.Add(buf, offset + 24), remoteAddr, 0, 16); uint remoteScope = (uint)Marshal.ReadInt32(buf, offset + 40); uint remotePort = (uint)Marshal.ReadInt32(buf, offset + 44); uint state = (uint)Marshal.ReadInt32(buf, offset + 48); offset += 56; if ((int)state != MIB_TCP_STATE_ESTAB) { continue; } if (HostPort(localPort) != SMB_PORT) { continue; } if (IsLoopbackV6(localAddr) || IsLoopbackV6(remoteAddr)) { continue; } Tcp445Conn c = new Tcp445Conn(); c.Af = AF_INET6; c.State = state; c.LocalAddrV6 = localAddr; c.RemoteAddrV6 = remoteAddr; c.LocalScopeId = localScope; c.RemoteScopeId = remoteScope; c.LocalPort = HostPort(localPort); c.RemotePort = HostPort(remotePort); c.LocalIp = new IPAddress(localAddr).ToString(); c.RemoteIp = new IPAddress(remoteAddr).ToString(); list.Add(c); } } finally { Marshal.FreeHGlobal(buf); } } private static IntPtr AllocRow(Tcp445Conn conn) { if (conn.Af == AF_INET6) { IntPtr p = Marshal.AllocHGlobal(52); Marshal.WriteInt32(p, 0, (int)conn.State); if (conn.LocalAddrV6 != null) { Marshal.Copy(conn.LocalAddrV6, 0, IntPtr.Add(p, 4), 16); } Marshal.WriteInt32(p, 20, (int)conn.LocalScopeId); int lp = ((conn.LocalPort & 0xFF) << 8) | ((conn.LocalPort >> 8) & 0xFF); Marshal.WriteInt32(p, 24, lp); if (conn.RemoteAddrV6 != null) { Marshal.Copy(conn.RemoteAddrV6, 0, IntPtr.Add(p, 28), 16); } Marshal.WriteInt32(p, 44, (int)conn.RemoteScopeId); int rp = ((conn.RemotePort & 0xFF) << 8) | ((conn.RemotePort >> 8) & 0xFF); Marshal.WriteInt32(p, 48, rp); return p; } IntPtr v4 = Marshal.AllocHGlobal(20); Marshal.WriteInt32(v4, 0, (int)conn.State); Marshal.WriteInt32(v4, 4, (int)conn.LocalAddrV4); int lp4 = ((conn.LocalPort & 0xFF) << 8) | ((conn.LocalPort >> 8) & 0xFF); Marshal.WriteInt32(v4, 8, lp4); Marshal.WriteInt32(v4, 12, (int)conn.RemoteAddrV4); int rp4 = ((conn.RemotePort & 0xFF) << 8) | ((conn.RemotePort >> 8) & 0xFF); Marshal.WriteInt32(v4, 16, rp4); return v4; } } } '@ Add-Type -TypeDefinition $code -Language CSharp -ErrorAction Stop } function Get-SMBeatTcp445ConnectionKey { param( [Parameter(Mandatory = $true)] $Conn ) '{0}|{1}|{2}|{3}|{4}' -f $Conn.Af, $Conn.LocalIp, $Conn.LocalPort, $Conn.RemoteIp, $Conn.RemotePort } function Get-SMBeatTcp445LiveConnections { $result = New-Object System.Collections.Generic.List[object] try { Initialize-SMBeatTcp445Native } catch { return @() } $rows = @([SMBeat.Native.Tcp445]::ListEstablished()) foreach ($row in $rows) { [void][SMBeat.Native.Tcp445]::Enable($row) $ok = [SMBeat.Native.Tcp445]::TryRead($row) if (-not $ok) { continue } $result.Add([PSCustomObject]@{ Key = (Get-SMBeatTcp445ConnectionKey -Conn $row) RemoteIp = [string]$row.RemoteIp BytesOut = [uint64]$row.BytesOut BytesIn = [uint64]$row.BytesIn }) | Out-Null } return $result.ToArray() } function New-SMBeatTcp445Tracker { @{} } function Update-SMBeatTcp445Tracker { param( [Parameter(Mandatory = $true)] [hashtable]$Tracker, [AllowEmptyCollection()] [object[]]$Live ) foreach ($key in @($Tracker.Keys)) { $Tracker[$key].Seen = $false } if ($null -eq $Live) { $Live = @() } foreach ($row in $Live) { $key = [string]$row.Key $out = [uint64]$row.BytesOut $inb = [uint64]$row.BytesIn if (-not $Tracker.ContainsKey($key)) { $Tracker[$key] = @{ RemoteIp = [string]$row.RemoteIp LastOut = $out LastIn = $inb PendingOut = [int64]0 PendingIn = [int64]0 Hits = 1 Seen = $true } continue } $e = $Tracker[$key] $e.Hits = [int]$e.Hits + 1 $e.PendingOut += (Get-SMBeatByteDelta -Current $out -Previous ([uint64]$e.LastOut)) $e.PendingIn += (Get-SMBeatByteDelta -Current $inb -Previous ([uint64]$e.LastIn)) $e.LastOut = $out $e.LastIn = $inb $e.Seen = $true $e.RemoteIp = [string]$row.RemoteIp } } function Get-SMBeatTcp445PendingDeltas { param( [Parameter(Mandatory = $true)] [hashtable]$Tracker, [switch]$Reset ) $byClient = @{} foreach ($key in @($Tracker.Keys)) { $e = $Tracker[$key] $ip = [string]$e.RemoteIp if ([string]::IsNullOrWhiteSpace($ip)) { continue } if (-not $byClient.ContainsKey($ip)) { $byClient[$ip] = @{ Sent = [int64]0 Received = [int64]0 } } $byClient[$ip].Sent += [int64]$e.PendingOut $byClient[$ip].Received += [int64]$e.PendingIn if ($Reset) { $e.PendingOut = [int64]0 $e.PendingIn = [int64]0 } } $list = New-Object System.Collections.Generic.List[object] foreach ($ip in $byClient.Keys) { $list.Add([PSCustomObject]@{ RemoteIp = $ip Sent = [int64]$byClient[$ip].Sent Received = [int64]$byClient[$ip].Received }) | Out-Null } return $list.ToArray() } function Complete-SMBeatTcp445ShortLived { param( [Parameter(Mandatory = $true)] [hashtable]$Tracker ) foreach ($key in @($Tracker.Keys)) { $e = $Tracker[$key] if ([bool]$e.Seen) { continue } $hits = 1 if ($e.Hits) { $hits = [int]$e.Hits } if ($hits -gt 1) { continue } $e.PendingOut += [int64]([uint64]$e.LastOut) $e.PendingIn += [int64]([uint64]$e.LastIn) } } function Remove-SMBeatTcp445Stale { param( [Parameter(Mandatory = $true)] [hashtable]$Tracker ) foreach ($key in @($Tracker.Keys)) { if (-not [bool]$Tracker[$key].Seen) { $Tracker.Remove($key) } } } function ConvertTo-SMBeatTcp445FlushRecords { param( [AllowEmptyCollection()] [object[]]$Deltas, [Parameter(Mandatory = $true)] [string]$Server, [Parameter(Mandatory = $true)] [datetime]$Utc, [Parameter(Mandatory = $true)] [int]$IntervalSec ) $records = New-Object System.Collections.Generic.List[object] $sumSent = [int64]0 $sumRecv = [int64]0 if ($null -eq $Deltas) { $Deltas = @() } foreach ($row in $Deltas) { $sumSent += [int64]$row.Sent $sumRecv += [int64]$row.Received $snap = New-SMBeatCounterSnapshot -Kind 'tcp445' -Name ([string]$row.RemoteIp) ` -ReadRaw 0 -WriteRaw 0 -SentRaw 0 -ReceivedRaw 0 -Client ([string]$row.RemoteIp) $records.Add((ConvertTo-SMBeatSampleRecord -Snapshot $snap -Server $Server -Utc $Utc -IntervalSec $IntervalSec ` -ReadDelta 0 -WriteDelta 0 -SentDelta ([int64]$row.Sent) -ReceivedDelta ([int64]$row.Received))) | Out-Null } $serverSnap = New-SMBeatCounterSnapshot -Kind 'tcp445' -Name '_listen445' ` -ReadRaw 0 -WriteRaw 0 -SentRaw 0 -ReceivedRaw 0 $records.Insert(0, (ConvertTo-SMBeatSampleRecord -Snapshot $serverSnap -Server $Server -Utc $Utc -IntervalSec $IntervalSec ` -ReadDelta 0 -WriteDelta 0 -SentDelta $sumSent -ReceivedDelta $sumRecv)) return $records.ToArray() } |