support multi key pattern matching (#25)

cmd/chatlog/log.go

@@ -22,6 +22,8 @@ func initLog(cmd *cobra.Command, args []string) {
 	if Debug {
 		zerolog.SetGlobalLevel(zerolog.DebugLevel)
 	}
+
+	log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr, TimeFormat: time.RFC3339})
 }
 
 func initTuiLog(cmd *cobra.Command, args []string) {

internal/wechat/decrypt/validator.go

@@ -15,13 +15,17 @@ type Validator struct {
 }
 
 // NewValidator 创建一个仅用于验证的验证器
-func NewValidator(dataDir string, platform string, version int) (*Validator, error) {
+func NewValidator(platform string, version int, dataDir string) (*Validator, error) {
+	dbFile := GetSimpleDBFile(platform, version)
+	dbPath := filepath.Join(dataDir + "/" + dbFile)
+	return NewValidatorWithFile(platform, version, dbPath)
+}
+
+func NewValidatorWithFile(platform string, version int, dbPath string) (*Validator, error) {
 	decryptor, err := NewDecryptor(platform, version)
 	if err != nil {
 		return nil, err
 	}
-	dbFile := GetSimpleDBFile(platform, version)
-	dbPath := filepath.Join(dataDir + "/" + dbFile)
 	d, err := common.OpenDBFile(dbPath, decryptor.GetPageSize())
 	if err != nil {
 		return nil, err

internal/wechat/key/darwin/v3.go

@@ -19,12 +19,22 @@ const (
 	MaxWorkersV3 = 8
 )
 
+var V3KeyPatterns = []KeyPatternInfo{
+	{
+		Pattern: []byte{0x72, 0x74, 0x72, 0x65, 0x65, 0x5f, 0x69, 0x33, 0x32},
+		Offset:  24,
+	},
+}
+
 type V3Extractor struct {
-	validator *decrypt.Validator
+	validator   *decrypt.Validator
+	keyPatterns []KeyPatternInfo
 }
 
 func NewV3Extractor() *V3Extractor {
-	return &V3Extractor{}
+	return &V3Extractor{
+		keyPatterns: V3KeyPatterns,
+	}
 }
 
 func (e *V3Extractor) Extract(ctx context.Context, proc *model.Process) (string, error) {
@@ -127,7 +137,6 @@ func (e *V3Extractor) findMemory(ctx context.Context, pid uint32, memoryChannel
 
 // worker processes memory regions to find V3 version key
 func (e *V3Extractor) worker(ctx context.Context, memoryChannel <-chan []byte, resultChannel chan<- string) {
-	keyPattern := []byte{0x72, 0x74, 0x72, 0x65, 0x65, 0x5f, 0x69, 0x33, 0x32}
 	for {
 		select {
 		case <-ctx.Done():
@@ -137,51 +146,60 @@ func (e *V3Extractor) worker(ctx context.Context, memoryChannel <-chan []byte, r
 				return
 			}
 
-			index := len(memory)
+			if key, ok := e.SearchKey(ctx, memory); ok {
-
-			for {
 				select {
-				case <-ctx.Done():
+				case resultChannel <- key:
-					return // Exit if context cancelled
 				default:
 				}
-
-				log.Debug().Msgf("Searching for V3 key in memory region, size: %d bytes", len(memory))
-
-				// Find pattern from end to beginning
-				index = bytes.LastIndex(memory[:index], keyPattern)
-				if index == -1 {
-					break // No more matches found
-				}
-
-				log.Debug().Msgf("Found potential V3 key pattern in memory region, index: %d", index)
-
-				// For V3, the key is 32 bytes and starts right after the pattern
-				if index+24+32 > len(memory) {
-					index -= 1
-					continue
-				}
-
-				// Extract the key data, which is right after the pattern and 32 bytes long
-				keyOffset := index + 24
-				keyData := memory[keyOffset : keyOffset+32]
-
-				// Validate key against database header
-				if e.validator.Validate(keyData) {
-					select {
-					case resultChannel <- hex.EncodeToString(keyData):
-						log.Debug().Msg("Key found: " + hex.EncodeToString(keyData))
-						return
-					default:
-					}
-				}
-
-				index -= 1
 			}
 		}
 	}
 }
 
+func (e *V3Extractor) SearchKey(ctx context.Context, memory []byte) (string, bool) {
+	for _, keyPattern := range e.keyPatterns {
+		index := len(memory)
+
+		for {
+			select {
+			case <-ctx.Done():
+				return "", false
+			default:
+			}
+
+			// Find pattern from end to beginning
+			index = bytes.LastIndex(memory[:index], keyPattern.Pattern)
+			if index == -1 {
+				break // No more matches found
+			}
+
+			// Check if we have enough space for the key
+			keyOffset := index + keyPattern.Offset
+			if keyOffset < 0 || keyOffset+32 > len(memory) {
+				index -= 1
+				continue
+			}
+
+			// Extract the key data, which is 32 bytes long
+			keyData := memory[keyOffset : keyOffset+32]
+
+			// Validate key against database header
+			if e.validator.Validate(keyData) {
+				log.Debug().
+					Str("pattern", hex.EncodeToString(keyPattern.Pattern)).
+					Int("offset", keyPattern.Offset).
+					Str("key", hex.EncodeToString(keyData)).
+					Msg("Key found")
+				return hex.EncodeToString(keyData), true
+			}
+
+			index -= 1
+		}
+	}
+
+	return "", false
+}
+
 func (e *V3Extractor) SetValidate(validator *decrypt.Validator) {
 	e.validator = validator
 }

internal/wechat/key/darwin/v4.go

@@ -19,12 +19,26 @@ const (
 	MaxWorkers = 8
 )
 
+var V4KeyPatterns = []KeyPatternInfo{
+	{
+		Pattern: []byte{0x20, 0x66, 0x74, 0x73, 0x35, 0x28, 0x25, 0x00},
+		Offset:  16,
+	},
+	{
+		Pattern: []byte{0x20, 0x66, 0x74, 0x73, 0x35, 0x28, 0x25, 0x00},
+		Offset:  -80,
+	},
+}
+
 type V4Extractor struct {
-	validator *decrypt.Validator
+	validator   *decrypt.Validator
+	keyPatterns []KeyPatternInfo
 }
 
 func NewV4Extractor() *V4Extractor {
-	return &V4Extractor{}
+	return &V4Extractor{
+		keyPatterns: V4KeyPatterns,
+	}
 }
 
 func (e *V4Extractor) Extract(ctx context.Context, proc *model.Process) (string, error) {
@@ -127,8 +141,6 @@ func (e *V4Extractor) findMemory(ctx context.Context, pid uint32, memoryChannel
 
 // worker processes memory regions to find V4 version key
 func (e *V4Extractor) worker(ctx context.Context, memoryChannel <-chan []byte, resultChannel chan<- string) {
-	keyPattern := []byte{0x20, 0x66, 0x74, 0x73, 0x35, 0x28, 0x25, 0x00}
-
 	for {
 		select {
 		case <-ctx.Done():
@@ -138,47 +150,65 @@ func (e *V4Extractor) worker(ctx context.Context, memoryChannel <-chan []byte, r
 				return
 			}
 
-			index := len(memory)
+			if key, ok := e.SearchKey(ctx, memory); ok {
-
-			for {
 				select {
-				case <-ctx.Done():
+				case resultChannel <- key:
-					return // Exit if context cancelled
 				default:
 				}
-
-				// Find pattern from end to beginning
-				index = bytes.LastIndex(memory[:index], keyPattern)
-				if index == -1 {
-					break // No more matches found
-				}
-
-				// Check if we have enough space for the key
-				if index+16+32 > len(memory) {
-					index -= 1
-					continue
-				}
-
-				// Extract the key data, which is 16 bytes after the pattern and 32 bytes long
-				keyOffset := index + 16
-				keyData := memory[keyOffset : keyOffset+32]
-
-				// Validate key against database header
-				if e.validator.Validate(keyData) {
-					select {
-					case resultChannel <- hex.EncodeToString(keyData):
-						log.Debug().Msg("Key found: " + hex.EncodeToString(keyData))
-						return
-					default:
-					}
-				}
-
-				index -= 1
 			}
 		}
 	}
 }
 
+func (e *V4Extractor) SearchKey(ctx context.Context, memory []byte) (string, bool) {
+	for _, keyPattern := range e.keyPatterns {
+		index := len(memory)
+
+		for {
+			select {
+			case <-ctx.Done():
+				return "", false
+			default:
+			}
+
+			// Find pattern from end to beginning
+			index = bytes.LastIndex(memory[:index], keyPattern.Pattern)
+			if index == -1 {
+				break // No more matches found
+			}
+
+			// Check if we have enough space for the key
+			keyOffset := index + keyPattern.Offset
+			if keyOffset < 0 || keyOffset+32 > len(memory) {
+				index -= 1
+				continue
+			}
+
+			// Extract the key data, which is 16 bytes after the pattern and 32 bytes long
+			keyData := memory[keyOffset : keyOffset+32]
+
+			// Validate key against database header
+			if e.validator.Validate(keyData) {
+				log.Debug().
+					Str("pattern", hex.EncodeToString(keyPattern.Pattern)).
+					Int("offset", keyPattern.Offset).
+					Str("key", hex.EncodeToString(keyData)).
+					Msg("Key found")
+				return hex.EncodeToString(keyData), true
+			}
+
+			index -= 1
+		}
+	}
+
+	return "", false
+}
+
 func (e *V4Extractor) SetValidate(validator *decrypt.Validator) {
 	e.validator = validator
 }
+
+type KeyPatternInfo struct {
+	Pattern []byte
+	Offset  int
+}

internal/wechat/key/extractor.go

@@ -15,6 +15,9 @@ type Extractor interface {
 	// Extract 从进程中提取密钥
 	Extract(ctx context.Context, proc *model.Process) (string, error)
 
+	// SearchKey 在内存中搜索密钥
+	SearchKey(ctx context.Context, memory []byte) (string, bool)
+
 	SetValidate(validator *decrypt.Validator)
 }
 

internal/wechat/key/windows/v3.go

@@ -1,6 +1,8 @@
 package windows
 
 import (
+	"context"
+
 	"github.com/sjzar/chatlog/internal/wechat/decrypt"
 )
 
@@ -12,6 +14,11 @@ func NewV3Extractor() *V3Extractor {
 	return &V3Extractor{}
 }
 
+func (e *V3Extractor) SearchKey(ctx context.Context, memory []byte) (string, bool) {
+	// TODO : Implement the key search logic for V3
+	return "", false
+}
+
 func (e *V3Extractor) SetValidate(validator *decrypt.Validator) {
 	e.validator = validator
 }

internal/wechat/key/windows/v4.go

@@ -1,6 +1,8 @@
 package windows
 
 import (
+	"context"
+
 	"github.com/sjzar/chatlog/internal/wechat/decrypt"
 )
 
@@ -12,6 +14,11 @@ func NewV4Extractor() *V4Extractor {
 	return &V4Extractor{}
 }
 
+func (e *V4Extractor) SearchKey(ctx context.Context, memory []byte) (string, bool) {
+	// TODO : Implement the key search logic for V4
+	return "", false
+}
+
 func (e *V4Extractor) SetValidate(validator *decrypt.Validator) {
 	e.validator = validator
 }

internal/wechat/wechat.go

@@ -90,7 +90,7 @@ func (a *Account) GetKey(ctx context.Context) (string, error) {
 		return "", err
 	}
 
-	validator, err := decrypt.NewValidator(process.DataDir, process.Platform, process.Version)
+	validator, err := decrypt.NewValidator(process.Platform, process.Version, process.DataDir)
 	if err != nil {
 		return "", err
 	}