auto decrypt (#44)

This commit is contained in:
Sarv
2025-04-16 01:02:29 +08:00
committed by GitHub
parent f2aa923e99
commit 25d0b394e2
20 changed files with 2903 additions and 712 deletions

View File

@@ -9,23 +9,57 @@ import (
"strings"
"time"
"github.com/fsnotify/fsnotify"
_ "github.com/mattn/go-sqlite3"
"github.com/rs/zerolog/log"
"github.com/sjzar/chatlog/internal/errors"
"github.com/sjzar/chatlog/internal/model"
"github.com/sjzar/chatlog/pkg/util"
"github.com/sjzar/chatlog/internal/wechatdb/datasource/dbm"
)
const (
MessageFilePattern = "^MSG([0-9]?[0-9])?\\.db$"
ContactFilePattern = "^MicroMsg.db$"
ImageFilePattern = "^HardLinkImage\\.db$"
VideoFilePattern = "^HardLinkVideo\\.db$"
FileFilePattern = "^HardLinkFile\\.db$"
VoiceFilePattern = "^MediaMSG([0-9])?\\.db$"
Message = "message"
Contact = "contact"
Image = "image"
Video = "video"
File = "file"
Voice = "voice"
)
var Groups = []dbm.Group{
{
Name: Message,
Pattern: `^MSG([0-9]?[0-9])?\.db$`,
BlackList: []string{},
},
{
Name: Contact,
Pattern: `^MicroMsg.db$`,
BlackList: []string{},
},
{
Name: Image,
Pattern: `^HardLinkImage\.db$`,
BlackList: []string{},
},
{
Name: Video,
Pattern: `^HardLinkVideo\.db$`,
BlackList: []string{},
},
{
Name: File,
Pattern: `^HardLinkFile\.db$`,
BlackList: []string{},
},
{
Name: Voice,
Pattern: `^MediaMSG([0-9])?\.db$`,
BlackList: []string{},
},
}
// MessageDBInfo 保存消息数据库的信息
type MessageDBInfo struct {
FilePath string
@@ -36,67 +70,67 @@ type MessageDBInfo struct {
// DataSource 实现了 DataSource 接口
type DataSource struct {
// 消息数据库
messageFiles []MessageDBInfo
messageDbs map[string]*sql.DB
path string
dbm *dbm.DBManager
// 联系人数据库
contactDbFile string
contactDb *sql.DB
imageDb *sql.DB
videoDb *sql.DB
fileDb *sql.DB
voiceDb []*sql.DB
// 消息数据库信息
messageInfos []MessageDBInfo
}
// New 创建一个新的 WindowsV3DataSource
func New(path string) (*DataSource, error) {
ds := &DataSource{
messageFiles: make([]MessageDBInfo, 0),
messageDbs: make(map[string]*sql.DB),
voiceDb: make([]*sql.DB, 0),
path: path,
dbm: dbm.NewDBManager(path),
messageInfos: make([]MessageDBInfo, 0),
}
// 初始化消息数据库
if err := ds.initMessageDbs(path); err != nil {
for _, g := range Groups {
ds.dbm.AddGroup(g)
}
if err := ds.dbm.Start(); err != nil {
return nil, err
}
if err := ds.initMessageDbs(); err != nil {
return nil, errors.DBInitFailed(err)
}
// 初始化联系人数据库
if err := ds.initContactDb(path); err != nil {
return nil, errors.DBInitFailed(err)
}
if err := ds.initMediaDb(path); err != nil {
return nil, errors.DBInitFailed(err)
}
if err := ds.initVoiceDb(path); err != nil {
return nil, errors.DBInitFailed(err)
}
ds.dbm.AddCallback(Message, func(event fsnotify.Event) error {
if !event.Op.Has(fsnotify.Create) {
return nil
}
if err := ds.initMessageDbs(); err != nil {
log.Err(err).Msgf("Failed to reinitialize message DBs: %s", event.Name)
}
return nil
})
return ds, nil
}
// initMessageDbs 初始化消息数据库
func (ds *DataSource) initMessageDbs(path string) error {
// 查找所有消息数据库文件
files, err := util.FindFilesWithPatterns(path, MessageFilePattern, true)
if err != nil {
return errors.DBFileNotFound(path, MessageFilePattern, err)
func (ds *DataSource) SetCallback(name string, callback func(event fsnotify.Event) error) error {
if name == "chatroom" {
name = Contact
}
return ds.dbm.AddCallback(name, callback)
}
if len(files) == 0 {
return errors.DBFileNotFound(path, MessageFilePattern, nil)
// initMessageDbs 初始化消息数据库
func (ds *DataSource) initMessageDbs() error {
// 获取所有消息数据库文件路径
dbPaths, err := ds.dbm.GetDBPath(Message)
if err != nil {
return err
}
// 处理每个数据库文件
for _, filePath := range files {
// 连接数据库
db, err := sql.Open("sqlite3", filePath)
infos := make([]MessageDBInfo, 0)
for _, filePath := range dbPaths {
db, err := ds.dbm.OpenDB(filePath)
if err != nil {
log.Err(err).Msgf("连接数据库 %s 失败", filePath)
log.Err(err).Msgf("获取数据库 %s 失败", filePath)
continue
}
@@ -106,7 +140,6 @@ func (ds *DataSource) initMessageDbs(path string) error {
rows, err := db.Query("SELECT tableIndex, tableVersion, tableDesc FROM DBInfo")
if err != nil {
log.Err(err).Msgf("查询数据库 %s 的 DBInfo 表失败", filePath)
db.Close()
continue
}
@@ -133,7 +166,6 @@ func (ds *DataSource) initMessageDbs(path string) error {
rows, err = db.Query("SELECT UsrName FROM Name2ID")
if err != nil {
log.Err(err).Msgf("查询数据库 %s 的 Name2ID 表失败", filePath)
db.Close()
continue
}
@@ -150,123 +182,34 @@ func (ds *DataSource) initMessageDbs(path string) error {
rows.Close()
// 保存数据库信息
ds.messageFiles = append(ds.messageFiles, MessageDBInfo{
infos = append(infos, MessageDBInfo{
FilePath: filePath,
StartTime: startTime,
TalkerMap: talkerMap,
})
// 保存数据库连接
ds.messageDbs[filePath] = db
}
// 按照 StartTime 排序数据库文件
sort.Slice(ds.messageFiles, func(i, j int) bool {
return ds.messageFiles[i].StartTime.Before(ds.messageFiles[j].StartTime)
sort.Slice(infos, func(i, j int) bool {
return infos[i].StartTime.Before(infos[j].StartTime)
})
// 设置结束时间
for i := range ds.messageFiles {
if i == len(ds.messageFiles)-1 {
ds.messageFiles[i].EndTime = time.Now()
for i := range infos {
if i == len(infos)-1 {
infos[i].EndTime = time.Now()
} else {
ds.messageFiles[i].EndTime = ds.messageFiles[i+1].StartTime
infos[i].EndTime = infos[i+1].StartTime
}
}
return nil
}
// initContactDb 初始化联系人数据库
func (ds *DataSource) initContactDb(path string) error {
files, err := util.FindFilesWithPatterns(path, ContactFilePattern, true)
if err != nil {
return errors.DBFileNotFound(path, ContactFilePattern, err)
}
if len(files) == 0 {
return errors.DBFileNotFound(path, ContactFilePattern, nil)
}
ds.contactDbFile = files[0]
ds.contactDb, err = sql.Open("sqlite3", ds.contactDbFile)
if err != nil {
return errors.DBConnectFailed(ds.contactDbFile, err)
}
return nil
}
// initContactDb 初始化联系人数据库
func (ds *DataSource) initMediaDb(path string) error {
files, err := util.FindFilesWithPatterns(path, ImageFilePattern, true)
if err != nil {
return errors.DBFileNotFound(path, ImageFilePattern, err)
}
if len(files) == 0 {
return errors.DBFileNotFound(path, ImageFilePattern, nil)
}
ds.imageDb, err = sql.Open("sqlite3", files[0])
if err != nil {
return errors.DBConnectFailed(files[0], err)
}
files, err = util.FindFilesWithPatterns(path, VideoFilePattern, true)
if err != nil {
return errors.DBFileNotFound(path, VideoFilePattern, err)
}
if len(files) == 0 {
return errors.DBFileNotFound(path, VideoFilePattern, nil)
}
ds.videoDb, err = sql.Open("sqlite3", files[0])
if err != nil {
return errors.DBConnectFailed(files[0], err)
}
files, err = util.FindFilesWithPatterns(path, FileFilePattern, true)
if err != nil {
return errors.DBFileNotFound(path, FileFilePattern, err)
}
if len(files) == 0 {
return errors.DBFileNotFound(path, FileFilePattern, nil)
}
ds.fileDb, err = sql.Open("sqlite3", files[0])
if err != nil {
return errors.DBConnectFailed(files[0], err)
}
return nil
}
func (ds *DataSource) initVoiceDb(path string) error {
files, err := util.FindFilesWithPatterns(path, VoiceFilePattern, true)
if err != nil {
return errors.DBFileNotFound(path, VoiceFilePattern, err)
}
if len(files) == 0 {
return errors.DBFileNotFound(path, VoiceFilePattern, nil)
}
for _, file := range files {
db, err := sql.Open("sqlite3", file)
if err != nil {
return errors.DBConnectFailed(files[0], err)
}
ds.voiceDb = append(ds.voiceDb, db)
}
ds.messageInfos = infos
return nil
}
// getDBInfosForTimeRange 获取时间范围内的数据库信息
func (ds *DataSource) getDBInfosForTimeRange(startTime, endTime time.Time) []MessageDBInfo {
var dbs []MessageDBInfo
for _, info := range ds.messageFiles {
for _, info := range ds.messageInfos {
if info.StartTime.Before(endTime) && info.EndTime.After(startTime) {
dbs = append(dbs, info)
}
@@ -296,70 +239,19 @@ func (ds *DataSource) GetMessages(ctx context.Context, startTime, endTime time.T
return nil, err
}
db, ok := ds.messageDbs[dbInfo.FilePath]
if !ok {
db, err := ds.dbm.OpenDB(dbInfo.FilePath)
if err != nil {
log.Error().Msgf("数据库 %s 未打开", dbInfo.FilePath)
continue
}
// 构建查询条件
conditions := []string{"Sequence >= ? AND Sequence <= ?"}
args := []interface{}{startTime.Unix() * 1000, endTime.Unix() * 1000}
if len(talker) > 0 {
talkerID, ok := dbInfo.TalkerMap[talker]
if ok {
conditions = append(conditions, "TalkerId = ?")
args = append(args, talkerID)
} else {
conditions = append(conditions, "StrTalker = ?")
args = append(args, talker)
}
}
query := fmt.Sprintf(`
SELECT MsgSvrID, Sequence, CreateTime, StrTalker, IsSender,
Type, SubType, StrContent, CompressContent, BytesExtra
FROM MSG
WHERE %s
ORDER BY Sequence ASC
`, strings.Join(conditions, " AND "))
// 执行查询
rows, err := db.QueryContext(ctx, query, args...)
messages, err := ds.getMessagesFromDB(ctx, db, dbInfo, startTime, endTime, talker)
if err != nil {
log.Err(err).Msgf("查询数据库 %s 失败", dbInfo.FilePath)
log.Err(err).Msgf("数据库 %s 获取消息失败", dbInfo.FilePath)
continue
}
// 处理查询结果
for rows.Next() {
var msg model.MessageV3
var compressContent []byte
var bytesExtra []byte
err := rows.Scan(
&msg.MsgSvrID,
&msg.Sequence,
&msg.CreateTime,
&msg.StrTalker,
&msg.IsSender,
&msg.Type,
&msg.SubType,
&msg.StrContent,
&compressContent,
&bytesExtra,
)
if err != nil {
log.Err(err).Msg("扫描消息行失败")
continue
}
msg.CompressContent = compressContent
msg.BytesExtra = bytesExtra
totalMessages = append(totalMessages, msg.Wrap())
}
rows.Close()
totalMessages = append(totalMessages, messages...)
if limit+offset > 0 && len(totalMessages) >= limit+offset {
break
@@ -388,6 +280,11 @@ func (ds *DataSource) GetMessages(ctx context.Context, startTime, endTime time.T
// getMessagesSingleFile 从单个数据库文件获取消息
func (ds *DataSource) getMessagesSingleFile(ctx context.Context, dbInfo MessageDBInfo, startTime, endTime time.Time, talker string, limit, offset int) ([]*model.Message, error) {
db, err := ds.dbm.OpenDB(dbInfo.FilePath)
if err != nil {
return nil, errors.DBConnectFailed(dbInfo.FilePath, nil)
}
// 构建查询条件
conditions := []string{"Sequence >= ? AND Sequence <= ?"}
args := []interface{}{startTime.Unix() * 1000, endTime.Unix() * 1000}
@@ -419,7 +316,7 @@ func (ds *DataSource) getMessagesSingleFile(ctx context.Context, dbInfo MessageD
}
// 执行查询
rows, err := ds.messageDbs[dbInfo.FilePath].QueryContext(ctx, query, args...)
rows, err := db.QueryContext(ctx, query, args...)
if err != nil {
return nil, errors.QueryFailed(query, err)
}
@@ -453,6 +350,69 @@ func (ds *DataSource) getMessagesSingleFile(ctx context.Context, dbInfo MessageD
return totalMessages, nil
}
// getMessagesFromDB 从数据库获取消息
func (ds *DataSource) getMessagesFromDB(ctx context.Context, db *sql.DB, dbInfo MessageDBInfo, startTime, endTime time.Time, talker string) ([]*model.Message, error) {
// 构建查询条件
conditions := []string{"Sequence >= ? AND Sequence <= ?"}
args := []interface{}{startTime.Unix() * 1000, endTime.Unix() * 1000}
if len(talker) > 0 {
talkerID, ok := dbInfo.TalkerMap[talker]
if ok {
conditions = append(conditions, "TalkerId = ?")
args = append(args, talkerID)
} else {
conditions = append(conditions, "StrTalker = ?")
args = append(args, talker)
}
}
query := fmt.Sprintf(`
SELECT MsgSvrID, Sequence, CreateTime, StrTalker, IsSender,
Type, SubType, StrContent, CompressContent, BytesExtra
FROM MSG
WHERE %s
ORDER BY Sequence ASC
`, strings.Join(conditions, " AND "))
// 执行查询
rows, err := db.QueryContext(ctx, query, args...)
if err != nil {
return nil, errors.QueryFailed(query, err)
}
defer rows.Close()
// 处理查询结果
messages := []*model.Message{}
for rows.Next() {
var msg model.MessageV3
var compressContent []byte
var bytesExtra []byte
err := rows.Scan(
&msg.MsgSvrID,
&msg.Sequence,
&msg.CreateTime,
&msg.StrTalker,
&msg.IsSender,
&msg.Type,
&msg.SubType,
&msg.StrContent,
&compressContent,
&bytesExtra,
)
if err != nil {
return nil, errors.ScanRowFailed(err)
}
msg.CompressContent = compressContent
msg.BytesExtra = bytesExtra
messages = append(messages, msg.Wrap())
}
return messages, nil
}
// GetContacts 实现获取联系人信息的方法
func (ds *DataSource) GetContacts(ctx context.Context, key string, limit, offset int) ([]*model.Contact, error) {
var query string
@@ -478,7 +438,11 @@ func (ds *DataSource) GetContacts(ctx context.Context, key string, limit, offset
}
// 执行查询
rows, err := ds.contactDb.QueryContext(ctx, query, args...)
db, err := ds.dbm.GetDB(Contact)
if err != nil {
return nil, err
}
rows, err := db.QueryContext(ctx, query, args...)
if err != nil {
return nil, errors.QueryFailed(query, err)
}
@@ -516,7 +480,11 @@ func (ds *DataSource) GetChatRooms(ctx context.Context, key string, limit, offse
args = []interface{}{key}
// 执行查询
rows, err := ds.contactDb.QueryContext(ctx, query, args...)
db, err := ds.dbm.GetDB(Contact)
if err != nil {
return nil, err
}
rows, err := db.QueryContext(ctx, query, args...)
if err != nil {
return nil, errors.QueryFailed(query, err)
}
@@ -543,7 +511,7 @@ func (ds *DataSource) GetChatRooms(ctx context.Context, key string, limit, offse
contacts, err := ds.GetContacts(ctx, key, 1, 0)
if err == nil && len(contacts) > 0 && strings.HasSuffix(contacts[0].UserName, "@chatroom") {
// 再次尝试通过用户名查找群聊
rows, err := ds.contactDb.QueryContext(ctx,
rows, err := db.QueryContext(ctx,
`SELECT ChatRoomName, Reserved2, RoomData FROM ChatRoom WHERE ChatRoomName = ?`,
contacts[0].UserName)
@@ -593,7 +561,11 @@ func (ds *DataSource) GetChatRooms(ctx context.Context, key string, limit, offse
}
// 执行查询
rows, err := ds.contactDb.QueryContext(ctx, query, args...)
db, err := ds.dbm.GetDB(Contact)
if err != nil {
return nil, err
}
rows, err := db.QueryContext(ctx, query, args...)
if err != nil {
return nil, errors.QueryFailed(query, err)
}
@@ -647,7 +619,11 @@ func (ds *DataSource) GetSessions(ctx context.Context, key string, limit, offset
}
// 执行查询
rows, err := ds.contactDb.QueryContext(ctx, query, args...)
db, err := ds.dbm.GetDB(Contact)
if err != nil {
return nil, err
}
rows, err := db.QueryContext(ctx, query, args...)
if err != nil {
return nil, errors.QueryFailed(query, err)
}
@@ -688,25 +664,29 @@ func (ds *DataSource) GetMedia(ctx context.Context, _type string, key string) (*
return nil, errors.DecodeKeyFailed(err)
}
var db *sql.DB
var dbType string
var table1, table2 string
switch _type {
case "image":
db = ds.imageDb
dbType = Image
table1 = "HardLinkImageAttribute"
table2 = "HardLinkImageID"
case "video":
db = ds.videoDb
dbType = Video
table1 = "HardLinkVideoAttribute"
table2 = "HardLinkVideoID"
case "file":
db = ds.fileDb
dbType = File
table1 = "HardLinkFileAttribute"
table2 = "HardLinkFileID"
default:
return nil, errors.MediaTypeUnsupported(_type)
}
db, err := ds.dbm.GetDB(dbType)
if err != nil {
return nil, err
}
query := fmt.Sprintf(`
@@ -768,7 +748,12 @@ func (ds *DataSource) GetVoice(ctx context.Context, key string) (*model.Media, e
`
args := []interface{}{key}
for _, db := range ds.voiceDb {
dbs, err := ds.dbm.GetDBs(Voice)
if err != nil {
return nil, errors.DBConnectFailed("", err)
}
for _, db := range dbs {
rows, err := db.QueryContext(ctx, query, args...)
if err != nil {
return nil, errors.QueryFailed(query, err)
@@ -798,41 +783,5 @@ func (ds *DataSource) GetVoice(ctx context.Context, key string) (*model.Media, e
// Close 实现 DataSource 接口的 Close 方法
func (ds *DataSource) Close() error {
var errs []error
// 关闭消息数据库连接
for _, db := range ds.messageDbs {
if err := db.Close(); err != nil {
errs = append(errs, err)
}
}
// 关闭联系人数据库连接
if ds.contactDb != nil {
if err := ds.contactDb.Close(); err != nil {
errs = append(errs, err)
}
}
if ds.imageDb != nil {
if err := ds.imageDb.Close(); err != nil {
errs = append(errs, err)
}
}
if ds.videoDb != nil {
if err := ds.videoDb.Close(); err != nil {
errs = append(errs, err)
}
}
if ds.fileDb != nil {
if err := ds.fileDb.Close(); err != nil {
errs = append(errs, err)
}
}
if len(errs) > 0 {
return errors.DBCloseFailed(errs[0])
}
return nil
return ds.dbm.Close()
}