x

internal/wechat/key/windows/v4_windows.go

@@ -0,0 +1,213 @@
+package windows
+
+import (
+	"bytes"
+	"context"
+	"encoding/binary"
+	"encoding/hex"
+	"fmt"
+	"runtime"
+	"sync"
+	"unsafe"
+
+	"github.com/sirupsen/logrus"
+	"golang.org/x/sys/windows"
+
+	"github.com/sjzar/chatlog/internal/wechat/model"
+)
+
+const (
+	MEM_PRIVATE = 0x20000
+)
+
+func (e *V4Extractor) Extract(ctx context.Context, proc *model.Process) (string, error) {
+	if proc.Status == model.StatusOffline {
+		return "", ErrWeChatOffline
+	}
+
+	// Open process handle
+	handle, err := windows.OpenProcess(windows.PROCESS_VM_READ|windows.PROCESS_QUERY_INFORMATION, false, proc.PID)
+	if err != nil {
+		return "", fmt.Errorf("%w: %v", ErrOpenProcess, err)
+	}
+	defer windows.CloseHandle(handle)
+
+	// Create context to control all goroutines
+	searchCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	// Create channels for memory data and results
+	memoryChannel := make(chan []byte, 100)
+	resultChannel := make(chan string, 1)
+
+	// Determine number of worker goroutines
+	workerCount := runtime.NumCPU()
+	if workerCount < 2 {
+		workerCount = 2
+	}
+	if workerCount > MaxWorkers {
+		workerCount = MaxWorkers
+	}
+	logrus.Debug("Starting ", workerCount, " workers for V4 key search")
+
+	// Start consumer goroutines
+	var workerWaitGroup sync.WaitGroup
+	workerWaitGroup.Add(workerCount)
+	for index := 0; index < workerCount; index++ {
+		go func() {
+			defer workerWaitGroup.Done()
+			e.worker(searchCtx, handle, memoryChannel, resultChannel)
+		}()
+	}
+
+	// Start producer goroutine
+	var producerWaitGroup sync.WaitGroup
+	producerWaitGroup.Add(1)
+	go func() {
+		defer producerWaitGroup.Done()
+		defer close(memoryChannel) // Close channel when producer is done
+		err := e.findMemory(searchCtx, handle, memoryChannel)
+		if err != nil {
+			logrus.Error(err)
+		}
+	}()
+
+	// Wait for producer and consumers to complete
+	go func() {
+		producerWaitGroup.Wait()
+		workerWaitGroup.Wait()
+		close(resultChannel)
+	}()
+
+	// Wait for result
+	select {
+	case <-ctx.Done():
+		return "", ctx.Err()
+	case result, ok := <-resultChannel:
+		if ok && result != "" {
+			return result, nil
+		}
+	}
+
+	return "", ErrNoValidKey
+}
+
+// findMemoryV4 searches for writable memory regions for V4 version
+func (e *V4Extractor) findMemory(ctx context.Context, handle windows.Handle, memoryChannel chan<- []byte) error {
+	// Define search range
+	minAddr := uintptr(0x10000)    // Process space usually starts from 0x10000
+	maxAddr := uintptr(0x7FFFFFFF) // 32-bit process space limit
+
+	if runtime.GOARCH == "amd64" {
+		maxAddr = uintptr(0x7FFFFFFFFFFF) // 64-bit process space limit
+	}
+	logrus.Debug("Scanning memory regions from 0x", fmt.Sprintf("%X", minAddr), " to 0x", fmt.Sprintf("%X", maxAddr))
+
+	currentAddr := minAddr
+
+	for currentAddr < maxAddr {
+		var memInfo windows.MemoryBasicInformation
+		err := windows.VirtualQueryEx(handle, currentAddr, &memInfo, unsafe.Sizeof(memInfo))
+		if err != nil {
+			break
+		}
+
+		// Skip small memory regions
+		if memInfo.RegionSize < 1024*1024 {
+			currentAddr += uintptr(memInfo.RegionSize)
+			continue
+		}
+
+		// Check if memory region is readable and private
+		if memInfo.State == windows.MEM_COMMIT && (memInfo.Protect&windows.PAGE_READWRITE) != 0 && memInfo.Type == MEM_PRIVATE {
+			// Calculate region size, ensure it doesn't exceed limit
+			regionSize := uintptr(memInfo.RegionSize)
+			if currentAddr+regionSize > maxAddr {
+				regionSize = maxAddr - currentAddr
+			}
+
+			// Read memory region
+			memory := make([]byte, regionSize)
+			if err = windows.ReadProcessMemory(handle, currentAddr, &memory[0], regionSize, nil); err == nil {
+				select {
+				case memoryChannel <- memory:
+					logrus.Debug("Sent memory region for analysis, size: ", regionSize, " bytes")
+				case <-ctx.Done():
+					return nil
+				}
+			}
+		}
+
+		// Move to next memory region
+		currentAddr = uintptr(memInfo.BaseAddress) + uintptr(memInfo.RegionSize)
+	}
+
+	return nil
+}
+
+// workerV4 processes memory regions to find V4 version key
+func (e *V4Extractor) worker(ctx context.Context, handle windows.Handle, memoryChannel <-chan []byte, resultChannel chan<- string) {
+	// Define search pattern for V4
+	keyPattern := []byte{
+		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+		0x20, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+		0x2F, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+	}
+	ptrSize := 8
+	littleEndianFunc := binary.LittleEndian.Uint64
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case memory, ok := <-memoryChannel:
+			if !ok {
+				return
+			}
+
+			index := len(memory)
+			for {
+				select {
+				case <-ctx.Done():
+					return // Exit if context cancelled
+				default:
+				}
+
+				// Find pattern from end to beginning
+				index = bytes.LastIndex(memory[:index], keyPattern)
+				if index == -1 || index-ptrSize < 0 {
+					break
+				}
+
+				// Extract and validate pointer value
+				ptrValue := littleEndianFunc(memory[index-ptrSize : index])
+				if ptrValue > 0x10000 && ptrValue < 0x7FFFFFFFFFFF {
+					if key := e.validateKey(handle, ptrValue); key != "" {
+						select {
+						case resultChannel <- key:
+							logrus.Debug("Valid key found for V4 database")
+							return
+						default:
+						}
+					}
+				}
+				index -= 1 // Continue searching from previous position
+			}
+		}
+	}
+}
+
+// validateKey validates a single key candidate
+func (e *V4Extractor) validateKey(handle windows.Handle, addr uint64) string {
+	keyData := make([]byte, 0x20) // 32-byte key
+	if err := windows.ReadProcessMemory(handle, uintptr(addr), &keyData[0], uintptr(len(keyData)), nil); err != nil {
+		return ""
+	}
+
+	// Validate key against database header
+	if e.validator.Validate(keyData) {
+		return hex.EncodeToString(keyData)
+	}
+
+	return ""
+}