first commit
This commit is contained in:
@@ -0,0 +1,619 @@
|
||||
package session
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net"
|
||||
"os"
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
"unicode/utf16"
|
||||
"unsafe"
|
||||
|
||||
"monitoring-agent-win/api"
|
||||
"monitoring-agent-win/logging"
|
||||
|
||||
"golang.org/x/sys/windows"
|
||||
)
|
||||
|
||||
// Event types
|
||||
const (
|
||||
EventLogon = "Logon"
|
||||
EventLogoff = "Logoff"
|
||||
EventLock = "Lock"
|
||||
EventUnlock = "Unlock"
|
||||
)
|
||||
|
||||
// WTS constants
|
||||
const (
|
||||
WTS_CURRENT_SERVER_HANDLE = 0
|
||||
WTS_SESSION_LOGON = 0x5
|
||||
WTS_SESSION_LOGOFF = 0x6
|
||||
WTS_SESSION_LOCK = 0x7
|
||||
WTS_SESSION_UNLOCK = 0x8
|
||||
WTS_USERNAME = 5
|
||||
NOTIFY_FOR_ALL_SESSIONS = 1
|
||||
WM_WTSSESSION_CHANGE = 0x02B1
|
||||
)
|
||||
|
||||
// Windows constants
|
||||
const (
|
||||
PM_REMOVE = 0x0001
|
||||
)
|
||||
|
||||
// Windows types
|
||||
type (
|
||||
HWND uintptr
|
||||
HINSTANCE uintptr
|
||||
LPARAM uintptr
|
||||
WPARAM uintptr
|
||||
LRESULT uintptr
|
||||
HANDLE uintptr
|
||||
HMODULE uintptr
|
||||
WNDPROC uintptr
|
||||
)
|
||||
|
||||
// Point represents a point structure
|
||||
type Point struct {
|
||||
X int32
|
||||
Y int32
|
||||
}
|
||||
|
||||
// WNDCLASSEX represents a window class
|
||||
type WNDCLASSEX struct {
|
||||
Size uint32
|
||||
Style uint32
|
||||
WndProc uintptr
|
||||
ClsExtra int32
|
||||
WndExtra int32
|
||||
Instance HINSTANCE
|
||||
Icon uintptr
|
||||
Cursor uintptr
|
||||
Background uintptr
|
||||
MenuName *uint16
|
||||
ClassName *uint16
|
||||
IconSm uintptr
|
||||
}
|
||||
|
||||
// MSG represents a message
|
||||
type MSG struct {
|
||||
Hwnd HWND
|
||||
Message uint32
|
||||
WParam WPARAM
|
||||
LParam LPARAM
|
||||
Time uint32
|
||||
Pt Point
|
||||
}
|
||||
|
||||
// Monitor represents a session monitor
|
||||
type Monitor struct {
|
||||
apiClient *api.Client
|
||||
csvLogger *logging.CSVLogger
|
||||
logger *logging.Logger
|
||||
timezone string
|
||||
lastUserMap map[uint32]string
|
||||
mu sync.RWMutex
|
||||
stopChan chan struct{}
|
||||
hwnd HWND
|
||||
}
|
||||
|
||||
// NewMonitor creates a new session monitor
|
||||
func NewMonitor(apiClient *api.Client, csvLogger *logging.CSVLogger, logger *logging.Logger, timezone string) *Monitor {
|
||||
return &Monitor{
|
||||
apiClient: apiClient,
|
||||
csvLogger: csvLogger,
|
||||
logger: logger,
|
||||
timezone: timezone,
|
||||
lastUserMap: make(map[uint32]string),
|
||||
stopChan: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Start starts the session monitor
|
||||
func (m *Monitor) Start() error {
|
||||
// Load required DLLs
|
||||
user32, err := windows.LoadDLL("user32.dll")
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to load user32.dll: %v", err)
|
||||
}
|
||||
defer user32.Release()
|
||||
|
||||
kernel32, err := windows.LoadDLL("kernel32.dll")
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to load kernel32.dll: %v", err)
|
||||
}
|
||||
defer kernel32.Release()
|
||||
|
||||
// Get procedures
|
||||
registerClassEx, err := user32.FindProc("RegisterClassExW")
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to find RegisterClassExW: %v", err)
|
||||
}
|
||||
|
||||
createWindowEx, err := user32.FindProc("CreateWindowExW")
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to find CreateWindowExW: %v", err)
|
||||
}
|
||||
|
||||
getModuleHandle, err := kernel32.FindProc("GetModuleHandleW")
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to find GetModuleHandleW: %v", err)
|
||||
}
|
||||
|
||||
// Register window class
|
||||
className, err := syscall.UTF16PtrFromString("SessionMonitorWindow")
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to convert class name: %v", err)
|
||||
}
|
||||
|
||||
// Get module handle
|
||||
hInstance, _, _ := getModuleHandle.Call(0)
|
||||
if hInstance == 0 {
|
||||
return fmt.Errorf("failed to get module handle")
|
||||
}
|
||||
|
||||
// Define window procedure callback
|
||||
wndProcCallback := syscall.NewCallback(m.wndProc)
|
||||
|
||||
// Register window class
|
||||
wndClass := WNDCLASSEX{
|
||||
Size: uint32(unsafe.Sizeof(WNDCLASSEX{})),
|
||||
WndProc: wndProcCallback,
|
||||
Instance: HINSTANCE(hInstance),
|
||||
ClassName: className,
|
||||
}
|
||||
|
||||
_, _, err = registerClassEx.Call(uintptr(unsafe.Pointer(&wndClass)))
|
||||
if err != syscall.Errno(0) {
|
||||
return fmt.Errorf("failed to register window class: %v", err)
|
||||
}
|
||||
|
||||
// Create window
|
||||
windowName, err := syscall.UTF16PtrFromString("Session Monitor")
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to convert window name: %v", err)
|
||||
}
|
||||
|
||||
hwnd, _, err := createWindowEx.Call(
|
||||
0,
|
||||
uintptr(unsafe.Pointer(className)),
|
||||
uintptr(unsafe.Pointer(windowName)),
|
||||
0,
|
||||
0, 0, 0, 0,
|
||||
0, 0, hInstance, 0,
|
||||
)
|
||||
if hwnd == 0 {
|
||||
return fmt.Errorf("failed to create window: %v", err)
|
||||
}
|
||||
|
||||
m.hwnd = HWND(hwnd)
|
||||
|
||||
// Register for session notifications
|
||||
if err := wtsRegisterSessionNotification(m.hwnd, NOTIFY_FOR_ALL_SESSIONS); err != nil {
|
||||
destroyWindow(m.hwnd)
|
||||
return fmt.Errorf("failed to register for session notifications: %v", err)
|
||||
}
|
||||
|
||||
// Start message loop in a goroutine
|
||||
go m.messageLoop()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Stop stops the session monitor
|
||||
func (m *Monitor) Stop() {
|
||||
close(m.stopChan)
|
||||
|
||||
// Unregister session notifications and destroy window
|
||||
if m.hwnd != 0 {
|
||||
wtsUnRegisterSessionNotification(m.hwnd)
|
||||
destroyWindow(m.hwnd)
|
||||
m.hwnd = 0
|
||||
}
|
||||
}
|
||||
|
||||
// messageLoop processes window messages
|
||||
func (m *Monitor) messageLoop() {
|
||||
user32, err := windows.LoadDLL("user32.dll")
|
||||
if err != nil {
|
||||
m.logger.Error("Failed to load user32.dll: %v", err)
|
||||
return
|
||||
}
|
||||
defer user32.Release()
|
||||
|
||||
peekMessage, err := user32.FindProc("PeekMessageW")
|
||||
if err != nil {
|
||||
m.logger.Error("Failed to find PeekMessageW: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
translateMessage, err := user32.FindProc("TranslateMessage")
|
||||
if err != nil {
|
||||
m.logger.Error("Failed to find TranslateMessage: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
dispatchMessage, err := user32.FindProc("DispatchMessageW")
|
||||
if err != nil {
|
||||
m.logger.Error("Failed to find DispatchMessageW: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
var msg MSG
|
||||
for {
|
||||
select {
|
||||
case <-m.stopChan:
|
||||
return
|
||||
default:
|
||||
// Process messages
|
||||
ret, _, _ := peekMessage.Call(
|
||||
uintptr(unsafe.Pointer(&msg)),
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
PM_REMOVE,
|
||||
)
|
||||
|
||||
if ret != 0 {
|
||||
translateMessage.Call(uintptr(unsafe.Pointer(&msg)))
|
||||
dispatchMessage.Call(uintptr(unsafe.Pointer(&msg)))
|
||||
}
|
||||
|
||||
// Sleep to avoid high CPU usage
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// wndProc is the window procedure callback
|
||||
func (m *Monitor) wndProc(hwnd HWND, msg uint32, wParam, lParam uintptr) uintptr {
|
||||
switch msg {
|
||||
case WM_WTSSESSION_CHANGE:
|
||||
sessionID := uint32(lParam)
|
||||
eventType := uint32(wParam)
|
||||
|
||||
// Get session info
|
||||
username := m.getSessionUsername(sessionID)
|
||||
computerName := getComputerName()
|
||||
ipAddress := getDefaultIPv4()
|
||||
|
||||
// Log the event
|
||||
switch eventType {
|
||||
case WTS_SESSION_LOGON:
|
||||
m.logEvent(EventLogon, username, computerName, ipAddress)
|
||||
case WTS_SESSION_LOGOFF:
|
||||
m.logEvent(EventLogoff, username, computerName, ipAddress)
|
||||
case WTS_SESSION_LOCK:
|
||||
m.logEvent(EventLock, username, computerName, ipAddress)
|
||||
case WTS_SESSION_UNLOCK:
|
||||
m.logEvent(EventUnlock, username, computerName, ipAddress)
|
||||
}
|
||||
}
|
||||
|
||||
return defWindowProc(hwnd, msg, wParam, lParam)
|
||||
}
|
||||
|
||||
// defWindowProc calls the default window procedure
|
||||
func defWindowProc(hwnd HWND, msg uint32, wParam, lParam uintptr) uintptr {
|
||||
user32, err := windows.LoadDLL("user32.dll")
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
defer user32.Release()
|
||||
|
||||
proc, err := user32.FindProc("DefWindowProcW")
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
|
||||
ret, _, _ := proc.Call(
|
||||
uintptr(hwnd),
|
||||
uintptr(msg),
|
||||
wParam,
|
||||
lParam,
|
||||
)
|
||||
|
||||
return ret
|
||||
}
|
||||
|
||||
// destroyWindow destroys a window
|
||||
func destroyWindow(hwnd HWND) {
|
||||
user32, err := windows.LoadDLL("user32.dll")
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer user32.Release()
|
||||
|
||||
proc, err := user32.FindProc("DestroyWindow")
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
proc.Call(uintptr(hwnd))
|
||||
}
|
||||
|
||||
// logEvent logs an event to the CSV file and sends it to the API
|
||||
func (m *Monitor) logEvent(eventType, username, computerName, ipAddress string) {
|
||||
// Get current timestamp in the configured timezone
|
||||
var timestamp string
|
||||
if m.timezone != "" {
|
||||
// Load the configured timezone
|
||||
loc, err := time.LoadLocation(m.timezone)
|
||||
if err != nil {
|
||||
// If timezone loading fails, fall back to local time but log the error
|
||||
m.logger.Warning("Failed to load timezone '%s': %v - using local time", m.timezone, err)
|
||||
timestamp = time.Now().Format("2006-01-02T15:04:05-07:00")
|
||||
} else {
|
||||
// Use the configured timezone
|
||||
timestamp = time.Now().In(loc).Format("2006-01-02T15:04:05-07:00")
|
||||
}
|
||||
} else {
|
||||
// If no timezone configured, use local time
|
||||
timestamp = time.Now().Format("2006-01-02T15:04:05-07:00")
|
||||
}
|
||||
|
||||
// Log the event
|
||||
m.logger.Info("User: %s - Event: %s - Computer: %s - IP: %s", username, eventType, computerName, ipAddress)
|
||||
|
||||
// Send to API
|
||||
success, err := m.apiClient.SendEvent(eventType, username, computerName, ipAddress, timestamp, 0)
|
||||
if err != nil {
|
||||
m.logger.Error("Failed to send event: %v", err)
|
||||
}
|
||||
|
||||
// Log to CSV
|
||||
sendStatus := logging.SendStatusSuccess
|
||||
if !success {
|
||||
sendStatus = logging.SendStatusFailed
|
||||
}
|
||||
|
||||
if err := m.csvLogger.LogEvent(eventType, username, computerName, ipAddress, timestamp, sendStatus); err != nil {
|
||||
m.logger.Error("Failed to log event to CSV: %v", err)
|
||||
}
|
||||
|
||||
// If the current event was sent successfully, retry failed events
|
||||
if success {
|
||||
m.retryFailedEvents()
|
||||
}
|
||||
}
|
||||
|
||||
// retryFailedEvents retries sending failed events
|
||||
func (m *Monitor) retryFailedEvents() {
|
||||
// Get failed events
|
||||
failedEvents, err := m.csvLogger.GetFailedEvents()
|
||||
if err != nil {
|
||||
m.logger.Error("Failed to get failed events: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
if len(failedEvents) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// Try to send each failed event
|
||||
successCount := 0
|
||||
var successfulEvents []map[string]string
|
||||
for _, event := range failedEvents {
|
||||
// Send to API with retry flag
|
||||
success, err := m.apiClient.SendEvent(
|
||||
event["eventtype"],
|
||||
event["username"],
|
||||
event["hostname"],
|
||||
event["ipaddress"],
|
||||
event["timestamp"],
|
||||
1, // Retry flag
|
||||
)
|
||||
|
||||
if err != nil {
|
||||
m.logger.Error("Failed to retry event: %v", err)
|
||||
continue
|
||||
}
|
||||
|
||||
if success {
|
||||
// Log to main CSV with retry success status
|
||||
if err := m.csvLogger.LogEvent(
|
||||
event["eventtype"],
|
||||
event["username"],
|
||||
event["hostname"],
|
||||
event["ipaddress"],
|
||||
event["timestamp"],
|
||||
logging.SendStatusRetrySuccess,
|
||||
); err != nil {
|
||||
m.logger.Error("Failed to log retry success to CSV: %v", err)
|
||||
}
|
||||
successCount++
|
||||
successfulEvents = append(successfulEvents, event)
|
||||
}
|
||||
}
|
||||
|
||||
// Remove successful events from retry file
|
||||
if successCount > 0 {
|
||||
if successCount == len(failedEvents) {
|
||||
// If all retries were successful, clear the entire retry file (more efficient)
|
||||
if err := m.csvLogger.ClearRetryFile(); err != nil {
|
||||
m.logger.Error("Failed to clear retry file: %v", err)
|
||||
}
|
||||
} else {
|
||||
// If only some retries were successful, remove just the successful ones
|
||||
if err := m.csvLogger.RemoveSuccessfulRetries(successfulEvents); err != nil {
|
||||
m.logger.Error("Failed to remove successful retries: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// getSessionUsername gets the username for a session
|
||||
func (m *Monitor) getSessionUsername(sessionID uint32) string {
|
||||
var username string
|
||||
var size uint32
|
||||
|
||||
// Try to get the username from WTS
|
||||
if wtsQuerySessionInformation(WTS_CURRENT_SERVER_HANDLE, sessionID, WTS_USERNAME, &username, &size) {
|
||||
if username != "" {
|
||||
// Cache the username
|
||||
m.mu.Lock()
|
||||
m.lastUserMap[sessionID] = username
|
||||
m.mu.Unlock()
|
||||
return username
|
||||
}
|
||||
}
|
||||
|
||||
// If we couldn't get the username, try to use the cached one
|
||||
m.mu.RLock()
|
||||
cachedUsername, ok := m.lastUserMap[sessionID]
|
||||
m.mu.RUnlock()
|
||||
|
||||
if ok {
|
||||
return cachedUsername
|
||||
}
|
||||
|
||||
// If all else fails, return "Unknown"
|
||||
return "Unknown"
|
||||
}
|
||||
|
||||
// getComputerName gets the computer name
|
||||
func getComputerName() string {
|
||||
hostname, err := os.Hostname()
|
||||
if err != nil {
|
||||
return "Unknown"
|
||||
}
|
||||
return hostname
|
||||
}
|
||||
|
||||
// getDefaultIPv4 gets the default IPv4 address
|
||||
func getDefaultIPv4() string {
|
||||
conn, err := net.Dial("udp", "8.8.8.8:80")
|
||||
if err != nil {
|
||||
return "Unknown"
|
||||
}
|
||||
defer conn.Close()
|
||||
|
||||
localAddr := conn.LocalAddr().(*net.UDPAddr)
|
||||
return localAddr.IP.String()
|
||||
}
|
||||
|
||||
// wtsRegisterSessionNotification registers for session notifications
|
||||
func wtsRegisterSessionNotification(hwnd HWND, dwFlags uint32) error {
|
||||
wtsapi32, err := windows.LoadDLL("wtsapi32.dll")
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to load wtsapi32.dll: %v", err)
|
||||
}
|
||||
defer wtsapi32.Release()
|
||||
|
||||
proc, err := wtsapi32.FindProc("WTSRegisterSessionNotification")
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to find WTSRegisterSessionNotification: %v", err)
|
||||
}
|
||||
|
||||
r1, _, e1 := proc.Call(uintptr(hwnd), uintptr(dwFlags))
|
||||
if r1 == 0 {
|
||||
if e1 != nil && e1 != syscall.Errno(0) {
|
||||
return fmt.Errorf("WTSRegisterSessionNotification failed: %v", e1)
|
||||
}
|
||||
return fmt.Errorf("WTSRegisterSessionNotification failed")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// wtsUnRegisterSessionNotification unregisters from session notifications
|
||||
func wtsUnRegisterSessionNotification(hwnd HWND) error {
|
||||
wtsapi32, err := windows.LoadDLL("wtsapi32.dll")
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to load wtsapi32.dll: %v", err)
|
||||
}
|
||||
defer wtsapi32.Release()
|
||||
|
||||
proc, err := wtsapi32.FindProc("WTSUnRegisterSessionNotification")
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to find WTSUnRegisterSessionNotification: %v", err)
|
||||
}
|
||||
|
||||
r1, _, e1 := proc.Call(uintptr(hwnd))
|
||||
if r1 == 0 {
|
||||
if e1 != nil && e1 != syscall.Errno(0) {
|
||||
return fmt.Errorf("WTSUnRegisterSessionNotification failed: %v", e1)
|
||||
}
|
||||
return fmt.Errorf("WTSUnRegisterSessionNotification failed")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// wtsQuerySessionInformation queries session information
|
||||
func wtsQuerySessionInformation(hServer uintptr, sessionID uint32, infoClass uint32, ppBuffer *string, pBytesReturned *uint32) bool {
|
||||
wtsapi32, err := windows.LoadDLL("wtsapi32.dll")
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
defer wtsapi32.Release()
|
||||
|
||||
proc, err := wtsapi32.FindProc("WTSQuerySessionInformationW")
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
var buffer unsafe.Pointer
|
||||
r1, _, _ := proc.Call(
|
||||
uintptr(hServer),
|
||||
uintptr(sessionID),
|
||||
uintptr(infoClass),
|
||||
uintptr(unsafe.Pointer(&buffer)),
|
||||
uintptr(unsafe.Pointer(pBytesReturned)),
|
||||
)
|
||||
|
||||
if r1 != 0 && buffer != nil {
|
||||
// Convert buffer to string
|
||||
bufferSize := *pBytesReturned / 2 // Size in WCHARs
|
||||
if bufferSize > 0 {
|
||||
// Convert to Go string
|
||||
*ppBuffer = utf16PtrToString((*uint16)(buffer))
|
||||
}
|
||||
|
||||
// Free the buffer
|
||||
wtsFreeMemory(uintptr(buffer))
|
||||
return true
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
// wtsFreeMemory frees memory allocated by WTS functions
|
||||
func wtsFreeMemory(pMemory uintptr) {
|
||||
if pMemory == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
wtsapi32, err := windows.LoadDLL("wtsapi32.dll")
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer wtsapi32.Release()
|
||||
|
||||
proc, err := wtsapi32.FindProc("WTSFreeMemory")
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
proc.Call(pMemory)
|
||||
}
|
||||
|
||||
// utf16PtrToString converts a UTF16 pointer to a Go string
|
||||
func utf16PtrToString(p *uint16) string {
|
||||
if p == nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
// Find the length of the string
|
||||
n := 0
|
||||
for ptr := unsafe.Pointer(p); *(*uint16)(ptr) != 0; ptr = unsafe.Pointer(uintptr(ptr) + 2) {
|
||||
n++
|
||||
}
|
||||
|
||||
// Create a slice and copy the data
|
||||
s := make([]uint16, n)
|
||||
for i := 0; i < n; i++ {
|
||||
s[i] = *(*uint16)(unsafe.Pointer(uintptr(unsafe.Pointer(p)) + uintptr(i*2)))
|
||||
}
|
||||
|
||||
return string(utf16.Decode(s))
|
||||
}
|
||||
Reference in New Issue
Block a user