x
This commit is contained in:
Shen Junzheng
409
internal/chatlog/app.go
Normal file
409
internal/chatlog/app.go
Normal file
@@ -0,0 +1,409 @@
|
||||
package chatlog
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/sjzar/chatlog/internal/chatlog/ctx"
|
||||
"github.com/sjzar/chatlog/internal/ui/footer"
|
||||
"github.com/sjzar/chatlog/internal/ui/help"
|
||||
"github.com/sjzar/chatlog/internal/ui/infobar"
|
||||
"github.com/sjzar/chatlog/internal/ui/menu"
|
||||
|
||||
"github.com/gdamore/tcell/v2"
|
||||
"github.com/rivo/tview"
|
||||
)
|
||||
|
||||
const (
|
||||
RefreshInterval = 1000 * time.Millisecond
|
||||
)
|
||||
|
||||
type App struct {
|
||||
*tview.Application
|
||||
|
||||
ctx *ctx.Context
|
||||
m *Manager
|
||||
stopRefresh chan struct{}
|
||||
|
||||
// page
|
||||
mainPages *tview.Pages
|
||||
infoBar *infobar.InfoBar
|
||||
tabPages *tview.Pages
|
||||
footer *footer.Footer
|
||||
|
||||
// tab
|
||||
menu *menu.Menu
|
||||
help *help.Help
|
||||
activeTab int
|
||||
tabCount int
|
||||
}
|
||||
|
||||
func NewApp(ctx *ctx.Context, m *Manager) *App {
|
||||
app := &App{
|
||||
ctx: ctx,
|
||||
m: m,
|
||||
Application: tview.NewApplication(),
|
||||
mainPages: tview.NewPages(),
|
||||
infoBar: infobar.New(),
|
||||
tabPages: tview.NewPages(),
|
||||
footer: footer.New(),
|
||||
menu: menu.New("主菜单"),
|
||||
help: help.New(),
|
||||
}
|
||||
|
||||
app.initMenu()
|
||||
|
||||
return app
|
||||
}
|
||||
|
||||
func (a *App) Run() error {
|
||||
|
||||
flex := tview.NewFlex().
|
||||
SetDirection(tview.FlexRow).
|
||||
AddItem(a.infoBar, infobar.InfoBarViewHeight, 0, false).
|
||||
AddItem(a.tabPages, 0, 1, true).
|
||||
AddItem(a.footer, 1, 1, false)
|
||||
|
||||
a.mainPages.AddPage("main", flex, true, true)
|
||||
|
||||
a.tabPages.
|
||||
AddPage("0", a.menu, true, true).
|
||||
AddPage("1", a.help, true, false)
|
||||
a.tabCount = 2
|
||||
|
||||
a.SetInputCapture(a.inputCapture)
|
||||
|
||||
go a.refresh()
|
||||
|
||||
if err := a.SetRoot(a.mainPages, true).EnableMouse(false).Run(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (a *App) Stop() {
|
||||
// 添加一个通道用于停止刷新 goroutine
|
||||
if a.stopRefresh != nil {
|
||||
close(a.stopRefresh)
|
||||
}
|
||||
a.Application.Stop()
|
||||
}
|
||||
|
||||
func (a *App) switchTab(step int) {
|
||||
index := (a.activeTab + step) % a.tabCount
|
||||
if index < 0 {
|
||||
index = a.tabCount - 1
|
||||
}
|
||||
a.activeTab = index
|
||||
a.tabPages.SwitchToPage(fmt.Sprint(a.activeTab))
|
||||
}
|
||||
|
||||
func (a *App) refresh() {
|
||||
tick := time.NewTicker(RefreshInterval)
|
||||
defer tick.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-a.stopRefresh:
|
||||
return
|
||||
case <-tick.C:
|
||||
a.infoBar.UpdateAccount(a.ctx.Account)
|
||||
a.infoBar.UpdateBasicInfo(a.ctx.PID, a.ctx.Version, a.ctx.ExePath)
|
||||
a.infoBar.UpdateStatus(a.ctx.Status)
|
||||
a.infoBar.UpdateDataKey(a.ctx.DataKey)
|
||||
a.infoBar.UpdateDataUsageDir(a.ctx.DataUsage, a.ctx.DataDir)
|
||||
a.infoBar.UpdateWorkUsageDir(a.ctx.WorkUsage, a.ctx.WorkDir)
|
||||
if a.ctx.HTTPEnabled {
|
||||
a.infoBar.UpdateHTTPServer(fmt.Sprintf("[green][已启动][white] [%s]", a.ctx.HTTPAddr))
|
||||
} else {
|
||||
a.infoBar.UpdateHTTPServer("[未启动]")
|
||||
}
|
||||
|
||||
a.Draw()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (a *App) inputCapture(event *tcell.EventKey) *tcell.EventKey {
|
||||
|
||||
// 如果当前页面不是主页面，ESC 键返回主页面
|
||||
if a.mainPages.HasPage("submenu") && event.Key() == tcell.KeyEscape {
|
||||
a.mainPages.RemovePage("submenu")
|
||||
a.mainPages.SwitchToPage("main")
|
||||
return nil
|
||||
}
|
||||
|
||||
if a.tabPages.HasFocus() {
|
||||
switch event.Key() {
|
||||
case tcell.KeyLeft:
|
||||
a.switchTab(-1)
|
||||
return nil
|
||||
case tcell.KeyRight:
|
||||
a.switchTab(1)
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
switch event.Key() {
|
||||
case tcell.KeyCtrlC:
|
||||
a.Stop()
|
||||
}
|
||||
|
||||
return event
|
||||
}
|
||||
|
||||
func (a *App) initMenu() {
|
||||
getDataKey := &menu.Item{
|
||||
Index: 2,
|
||||
Name: "获取数据密钥",
|
||||
Description: "从进程获取数据密钥",
|
||||
Selected: func(i *menu.Item) {
|
||||
if err := a.m.GetDataKey(); err != nil {
|
||||
a.showError(err)
|
||||
return
|
||||
}
|
||||
a.showInfo("获取数据密钥成功")
|
||||
},
|
||||
}
|
||||
|
||||
decryptData := &menu.Item{
|
||||
Index: 3,
|
||||
Name: "解密数据",
|
||||
Description: "解密数据文件",
|
||||
Selected: func(i *menu.Item) {
|
||||
// 创建一个没有按钮的模态框，显示"解密中..."
|
||||
modal := tview.NewModal().
|
||||
SetText("解密中...")
|
||||
|
||||
a.mainPages.AddPage("modal", modal, true, true)
|
||||
a.SetFocus(modal)
|
||||
|
||||
// 在后台执行解密操作
|
||||
go func() {
|
||||
// 执行解密
|
||||
err := a.m.DecryptDBFiles()
|
||||
|
||||
// 在主线程中更新UI
|
||||
a.QueueUpdateDraw(func() {
|
||||
if err != nil {
|
||||
// 解密失败
|
||||
modal.SetText("解密失败: " + err.Error())
|
||||
} else {
|
||||
// 解密成功
|
||||
modal.SetText("解密数据成功")
|
||||
}
|
||||
|
||||
// 添加确认按钮
|
||||
modal.AddButtons([]string{"OK"})
|
||||
modal.SetDoneFunc(func(buttonIndex int, buttonLabel string) {
|
||||
a.mainPages.RemovePage("modal")
|
||||
})
|
||||
a.SetFocus(modal)
|
||||
})
|
||||
}()
|
||||
},
|
||||
}
|
||||
|
||||
httpServer := &menu.Item{
|
||||
Index: 4,
|
||||
Name: "启动 HTTP 服务",
|
||||
Description: "启动本地 HTTP & MCP 服务器",
|
||||
Selected: func(i *menu.Item) {
|
||||
modal := tview.NewModal()
|
||||
|
||||
// 根据当前服务状态执行不同操作
|
||||
if !a.ctx.HTTPEnabled {
|
||||
// HTTP 服务未启动，启动服务
|
||||
modal.SetText("正在启动 HTTP 服务...")
|
||||
a.mainPages.AddPage("modal", modal, true, true)
|
||||
a.SetFocus(modal)
|
||||
|
||||
// 在后台启动服务
|
||||
go func() {
|
||||
err := a.m.StartService()
|
||||
|
||||
// 在主线程中更新UI
|
||||
a.QueueUpdateDraw(func() {
|
||||
if err != nil {
|
||||
// 启动失败
|
||||
modal.SetText("启动 HTTP 服务失败: " + err.Error())
|
||||
} else {
|
||||
// 启动成功
|
||||
modal.SetText("已启动 HTTP 服务")
|
||||
// 更改菜单项名称
|
||||
i.Name = "停止 HTTP 服务"
|
||||
i.Description = "停止本地 HTTP & MCP 服务器"
|
||||
}
|
||||
|
||||
// 添加确认按钮
|
||||
modal.AddButtons([]string{"OK"})
|
||||
modal.SetDoneFunc(func(buttonIndex int, buttonLabel string) {
|
||||
a.mainPages.RemovePage("modal")
|
||||
})
|
||||
a.SetFocus(modal)
|
||||
})
|
||||
}()
|
||||
} else {
|
||||
// HTTP 服务已启动，停止服务
|
||||
modal.SetText("正在停止 HTTP 服务...")
|
||||
a.mainPages.AddPage("modal", modal, true, true)
|
||||
a.SetFocus(modal)
|
||||
|
||||
// 在后台停止服务
|
||||
go func() {
|
||||
err := a.m.StopService()
|
||||
|
||||
// 在主线程中更新UI
|
||||
a.QueueUpdateDraw(func() {
|
||||
if err != nil {
|
||||
// 停止失败
|
||||
modal.SetText("停止 HTTP 服务失败: " + err.Error())
|
||||
} else {
|
||||
// 停止成功
|
||||
modal.SetText("已停止 HTTP 服务")
|
||||
// 更改菜单项名称
|
||||
i.Name = "启动 HTTP 服务"
|
||||
i.Description = "启动本地 HTTP 服务器"
|
||||
}
|
||||
|
||||
// 添加确认按钮
|
||||
modal.AddButtons([]string{"OK"})
|
||||
modal.SetDoneFunc(func(buttonIndex int, buttonLabel string) {
|
||||
a.mainPages.RemovePage("modal")
|
||||
})
|
||||
a.SetFocus(modal)
|
||||
})
|
||||
}()
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
setting := &menu.Item{
|
||||
Index: 5,
|
||||
Name: "设置",
|
||||
Description: "设置应用程序选项",
|
||||
Selected: a.settingSelected,
|
||||
}
|
||||
|
||||
a.menu.AddItem(setting)
|
||||
a.menu.AddItem(getDataKey)
|
||||
a.menu.AddItem(decryptData)
|
||||
a.menu.AddItem(httpServer)
|
||||
|
||||
a.menu.AddItem(&menu.Item{
|
||||
Index: 6,
|
||||
Name: "退出",
|
||||
Description: "退出程序",
|
||||
Selected: func(i *menu.Item) {
|
||||
a.Stop()
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// settingItem 表示一个设置项
|
||||
type settingItem struct {
|
||||
name string
|
||||
description string
|
||||
action func()
|
||||
}
|
||||
|
||||
func (a *App) settingSelected(i *menu.Item) {
|
||||
|
||||
settings := []settingItem{
|
||||
{
|
||||
name: "设置 HTTP 服务端口",
|
||||
description: "配置 HTTP 服务监听的端口",
|
||||
action: a.settingHTTPPort,
|
||||
},
|
||||
{
|
||||
name: "设置工作目录",
|
||||
description: "配置数据解密后的存储目录",
|
||||
action: a.settingWorkDir,
|
||||
},
|
||||
}
|
||||
|
||||
subMenu := menu.NewSubMenu("设置")
|
||||
for idx, setting := range settings {
|
||||
item := &menu.Item{
|
||||
Index: idx + 1,
|
||||
Name: setting.name,
|
||||
Description: setting.description,
|
||||
Selected: func(action func()) func(*menu.Item) {
|
||||
return func(*menu.Item) {
|
||||
action()
|
||||
}
|
||||
}(setting.action),
|
||||
}
|
||||
subMenu.AddItem(item)
|
||||
}
|
||||
|
||||
a.mainPages.AddPage("submenu", subMenu, true, true)
|
||||
a.SetFocus(subMenu)
|
||||
}
|
||||
|
||||
// settingHTTPPort 设置 HTTP 端口
|
||||
func (a *App) settingHTTPPort() {
|
||||
// 实现端口设置逻辑
|
||||
// 这里可以使用 tview.InputField 让用户输入端口
|
||||
form := tview.NewForm().
|
||||
AddInputField("端口", a.ctx.HTTPAddr, 20, nil, func(text string) {
|
||||
a.ctx.SetHTTPAddr(text)
|
||||
}).
|
||||
AddButton("保存", func() {
|
||||
a.mainPages.RemovePage("submenu2")
|
||||
a.showInfo("HTTP 端口已设置为 " + a.ctx.HTTPAddr)
|
||||
}).
|
||||
AddButton("取消", func() {
|
||||
a.mainPages.RemovePage("submenu2")
|
||||
})
|
||||
form.SetBorder(true).SetTitle("设置 HTTP 端口")
|
||||
|
||||
a.mainPages.AddPage("submenu2", form, true, true)
|
||||
a.SetFocus(form)
|
||||
}
|
||||
|
||||
// settingWorkDir 设置工作目录
|
||||
func (a *App) settingWorkDir() {
|
||||
// 实现工作目录设置逻辑
|
||||
form := tview.NewForm().
|
||||
AddInputField("工作目录", a.ctx.WorkDir, 40, nil, func(text string) {
|
||||
a.ctx.SetWorkDir(text)
|
||||
}).
|
||||
AddButton("保存", func() {
|
||||
a.mainPages.RemovePage("submenu2")
|
||||
a.showInfo("工作目录已设置为 " + a.ctx.WorkDir)
|
||||
}).
|
||||
AddButton("取消", func() {
|
||||
a.mainPages.RemovePage("submenu2")
|
||||
})
|
||||
form.SetBorder(true).SetTitle("设置工作目录")
|
||||
|
||||
a.mainPages.AddPage("submenu2", form, true, true)
|
||||
a.SetFocus(form)
|
||||
}
|
||||
|
||||
// showModal 显示一个模态对话框
|
||||
func (a *App) showModal(text string, buttons []string, doneFunc func(buttonIndex int, buttonLabel string)) {
|
||||
modal := tview.NewModal().
|
||||
SetText(text).
|
||||
AddButtons(buttons).
|
||||
SetDoneFunc(doneFunc)
|
||||
|
||||
a.mainPages.AddPage("modal", modal, true, true)
|
||||
a.SetFocus(modal)
|
||||
}
|
||||
|
||||
// showError 显示错误对话框
|
||||
func (a *App) showError(err error) {
|
||||
a.showModal(err.Error(), []string{"OK"}, func(buttonIndex int, buttonLabel string) {
|
||||
a.mainPages.RemovePage("modal")
|
||||
})
|
||||
}
|
||||
|
||||
// showInfo 显示信息对话框
|
||||
func (a *App) showInfo(text string) {
|
||||
a.showModal(text, []string{"OK"}, func(buttonIndex int, buttonLabel string) {
|
||||
a.mainPages.RemovePage("modal")
|
||||
})
|
||||
}
|
||||
60
internal/chatlog/conf/config.go
Normal file
60
internal/chatlog/conf/config.go
Normal file
@@ -0,0 +1,60 @@
|
||||
package conf
|
||||
|
||||
import "github.com/sjzar/chatlog/pkg/config"
|
||||
|
||||
type Config struct {
|
||||
ConfigDir string `mapstructure:"-"`
|
||||
LastAccount string `mapstructure:"last_account" json:"last_account"`
|
||||
History []ProcessConfig `mapstructure:"history" json:"history"`
|
||||
}
|
||||
|
||||
type ProcessConfig struct {
|
||||
Type string `mapstructure:"type" json:"type"`
|
||||
Version string `mapstructure:"version" json:"version"`
|
||||
MajorVersion int `mapstructure:"major_version" json:"major_version"`
|
||||
Account string `mapstructure:"account" json:"account"`
|
||||
DataKey string `mapstructure:"data_key" json:"data_key"`
|
||||
DataDir string `mapstructure:"data_dir" json:"data_dir"`
|
||||
WorkDir string `mapstructure:"work_dir" json:"work_dir"`
|
||||
HTTPEnabled bool `mapstructure:"http_enabled" json:"http_enabled"`
|
||||
HTTPAddr string `mapstructure:"http_addr" json:"http_addr"`
|
||||
LastTime int64 `mapstructure:"last_time" json:"last_time"`
|
||||
Files []File `mapstructure:"files" json:"files"`
|
||||
}
|
||||
|
||||
type File struct {
|
||||
Path string `mapstructure:"path" json:"path"`
|
||||
ModifiedTime int64 `mapstructure:"modified_time" json:"modified_time"`
|
||||
Size int64 `mapstructure:"size" json:"size"`
|
||||
}
|
||||
|
||||
func (c *Config) ParseHistory() map[string]ProcessConfig {
|
||||
m := make(map[string]ProcessConfig)
|
||||
for _, v := range c.History {
|
||||
m[v.Account] = v
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
func (c *Config) UpdateHistory(account string, conf ProcessConfig) error {
|
||||
if c.History == nil {
|
||||
c.History = make([]ProcessConfig, 0)
|
||||
}
|
||||
if len(c.History) == 0 {
|
||||
c.History = append(c.History, conf)
|
||||
} else {
|
||||
isFind := false
|
||||
for i, v := range c.History {
|
||||
if v.Account == account {
|
||||
isFind = true
|
||||
c.History[i] = conf
|
||||
break
|
||||
}
|
||||
}
|
||||
if !isFind {
|
||||
c.History = append(c.History, conf)
|
||||
}
|
||||
}
|
||||
config.SetConfig("last_account", account)
|
||||
return config.SetConfig("history", c.History)
|
||||
}
|
||||
68
internal/chatlog/conf/service.go
Normal file
68
internal/chatlog/conf/service.go
Normal file
@@ -0,0 +1,68 @@
|
||||
package conf
|
||||
|
||||
import (
|
||||
"log"
|
||||
"os"
|
||||
"sync"
|
||||
|
||||
"github.com/sjzar/chatlog/pkg/config"
|
||||
)
|
||||
|
||||
const (
|
||||
ConfigName = "chatlog"
|
||||
ConfigType = "json"
|
||||
EnvConfigDir = "CHATLOG_DIR"
|
||||
)
|
||||
|
||||
// Service 配置服务
|
||||
type Service struct {
|
||||
configPath string
|
||||
config *Config
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// NewService 创建配置服务
|
||||
func NewService(configPath string) (*Service, error) {
|
||||
|
||||
service := &Service{
|
||||
configPath: configPath,
|
||||
}
|
||||
|
||||
if err := service.Load(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return service, nil
|
||||
}
|
||||
|
||||
// Load 加载配置
|
||||
func (s *Service) Load() error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
configPath := s.configPath
|
||||
if configPath == "" {
|
||||
configPath = os.Getenv(EnvConfigDir)
|
||||
}
|
||||
if err := config.Init(ConfigName, ConfigType, configPath); err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
|
||||
conf := &Config{}
|
||||
if err := config.Load(conf); err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
conf.ConfigDir = config.ConfigPath
|
||||
s.config = conf
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetConfig 获取配置副本
|
||||
func (s *Service) GetConfig() *Config {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
// 返回配置副本
|
||||
configCopy := *s.config
|
||||
return &configCopy
|
||||
}
|
||||
158
internal/chatlog/ctx/context.go
Normal file
158
internal/chatlog/ctx/context.go
Normal file
@@ -0,0 +1,158 @@
|
||||
package ctx
|
||||
|
||||
import (
|
||||
"sync"
|
||||
|
||||
"github.com/sjzar/chatlog/internal/chatlog/conf"
|
||||
"github.com/sjzar/chatlog/internal/wechat"
|
||||
"github.com/sjzar/chatlog/pkg/util"
|
||||
)
|
||||
|
||||
// Context is a context for a chatlog.
|
||||
// It is used to store information about the chatlog.
|
||||
type Context struct {
|
||||
conf *conf.Service
|
||||
mu sync.RWMutex
|
||||
|
||||
History map[string]conf.ProcessConfig
|
||||
|
||||
// 微信账号相关状态
|
||||
Account string
|
||||
Version string
|
||||
MajorVersion int
|
||||
DataKey string
|
||||
DataUsage string
|
||||
DataDir string
|
||||
|
||||
// 工作目录相关状态
|
||||
WorkUsage string
|
||||
WorkDir string
|
||||
|
||||
// HTTP服务相关状态
|
||||
HTTPEnabled bool
|
||||
HTTPAddr string
|
||||
|
||||
// 当前选中的微信实例
|
||||
Current *wechat.Info
|
||||
PID int
|
||||
ExePath string
|
||||
Status string
|
||||
|
||||
// 所有可用的微信实例
|
||||
WeChatInstances []*wechat.Info
|
||||
}
|
||||
|
||||
func New(conf *conf.Service) *Context {
|
||||
ctx := &Context{
|
||||
conf: conf,
|
||||
}
|
||||
|
||||
ctx.loadConfig()
|
||||
|
||||
return ctx
|
||||
}
|
||||
|
||||
func (c *Context) loadConfig() {
|
||||
conf := c.conf.GetConfig()
|
||||
c.History = conf.ParseHistory()
|
||||
c.SwitchHistory(conf.LastAccount)
|
||||
c.Refresh()
|
||||
}
|
||||
|
||||
func (c *Context) SwitchHistory(account string) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
history, ok := c.History[account]
|
||||
if ok {
|
||||
c.Account = history.Account
|
||||
c.Version = history.Version
|
||||
c.MajorVersion = history.MajorVersion
|
||||
c.DataKey = history.DataKey
|
||||
c.DataDir = history.DataDir
|
||||
c.WorkDir = history.WorkDir
|
||||
c.HTTPEnabled = history.HTTPEnabled
|
||||
c.HTTPAddr = history.HTTPAddr
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Context) SwitchCurrent(info *wechat.Info) {
|
||||
c.SwitchHistory(info.AccountName)
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.Current = info
|
||||
c.Refresh()
|
||||
|
||||
}
|
||||
func (c *Context) Refresh() {
|
||||
if c.Current != nil {
|
||||
c.Account = c.Current.AccountName
|
||||
c.Version = c.Current.Version.FileVersion
|
||||
c.MajorVersion = c.Current.Version.FileMajorVersion
|
||||
c.PID = int(c.Current.PID)
|
||||
c.ExePath = c.Current.ExePath
|
||||
c.Status = c.Current.Status
|
||||
if c.Current.Key != "" && c.Current.Key != c.DataKey {
|
||||
c.DataKey = c.Current.Key
|
||||
}
|
||||
if c.Current.DataDir != "" && c.Current.DataDir != c.DataDir {
|
||||
c.DataDir = c.Current.DataDir
|
||||
}
|
||||
}
|
||||
if c.DataUsage == "" && c.DataDir != "" {
|
||||
go func() {
|
||||
c.DataUsage = util.GetDirSize(c.DataDir)
|
||||
}()
|
||||
}
|
||||
if c.WorkUsage == "" && c.WorkDir != "" {
|
||||
go func() {
|
||||
c.WorkUsage = util.GetDirSize(c.WorkDir)
|
||||
}()
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Context) SetHTTPEnabled(enabled bool) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.HTTPEnabled = enabled
|
||||
c.UpdateConfig()
|
||||
}
|
||||
|
||||
func (c *Context) SetHTTPAddr(addr string) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.HTTPAddr = addr
|
||||
c.UpdateConfig()
|
||||
}
|
||||
|
||||
func (c *Context) SetWorkDir(dir string) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.WorkDir = dir
|
||||
c.UpdateConfig()
|
||||
c.Refresh()
|
||||
}
|
||||
|
||||
func (c *Context) SetDataDir(dir string) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.DataDir = dir
|
||||
c.UpdateConfig()
|
||||
c.Refresh()
|
||||
}
|
||||
|
||||
// 更新配置
|
||||
func (c *Context) UpdateConfig() {
|
||||
pconf := conf.ProcessConfig{
|
||||
Type: "wechat",
|
||||
Version: c.Version,
|
||||
MajorVersion: c.MajorVersion,
|
||||
Account: c.Account,
|
||||
DataKey: c.DataKey,
|
||||
DataDir: c.DataDir,
|
||||
WorkDir: c.WorkDir,
|
||||
HTTPEnabled: c.HTTPEnabled,
|
||||
HTTPAddr: c.HTTPAddr,
|
||||
}
|
||||
conf := c.conf.GetConfig()
|
||||
conf.UpdateHistory(c.Account, pconf)
|
||||
}
|
||||
77
internal/chatlog/database/service.go
Normal file
77
internal/chatlog/database/service.go
Normal file
@@ -0,0 +1,77 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/sjzar/chatlog/internal/chatlog/ctx"
|
||||
"github.com/sjzar/chatlog/internal/wechatdb"
|
||||
"github.com/sjzar/chatlog/pkg/model"
|
||||
)
|
||||
|
||||
type Service struct {
|
||||
ctx *ctx.Context
|
||||
db *wechatdb.DB
|
||||
}
|
||||
|
||||
func NewService(ctx *ctx.Context) *Service {
|
||||
return &Service{
|
||||
ctx: ctx,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) Start() error {
|
||||
db, err := wechatdb.New(s.ctx.WorkDir, s.ctx.MajorVersion)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s.db = db
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Service) Stop() error {
|
||||
if s.db != nil {
|
||||
s.db.Close()
|
||||
}
|
||||
s.db = nil
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetDB returns the underlying database
|
||||
func (s *Service) GetDB() *wechatdb.DB {
|
||||
return s.db
|
||||
}
|
||||
|
||||
// GetMessages retrieves messages based on criteria
|
||||
func (s *Service) GetMessages(start, end time.Time, talker string, limit, offset int) ([]*model.Message, error) {
|
||||
return s.db.GetMessages(start, end, talker, limit, offset)
|
||||
}
|
||||
|
||||
// GetContact retrieves contact information
|
||||
func (s *Service) GetContact(userName string) *model.Contact {
|
||||
return s.db.GetContact(userName)
|
||||
}
|
||||
|
||||
// ListContact retrieves all contacts
|
||||
func (s *Service) ListContact() (*wechatdb.ListContactResp, error) {
|
||||
return s.db.ListContact()
|
||||
}
|
||||
|
||||
// GetChatRoom retrieves chat room information
|
||||
func (s *Service) GetChatRoom(name string) *model.ChatRoom {
|
||||
return s.db.GetChatRoom(name)
|
||||
}
|
||||
|
||||
// ListChatRoom retrieves all chat rooms
|
||||
func (s *Service) ListChatRoom() (*wechatdb.ListChatRoomResp, error) {
|
||||
return s.db.ListChatRoom()
|
||||
}
|
||||
|
||||
// GetSession retrieves session information
|
||||
func (s *Service) GetSession(limit int) (*wechatdb.GetSessionResp, error) {
|
||||
return s.db.GetSession(limit)
|
||||
}
|
||||
|
||||
// Close closes the database connection
|
||||
func (s *Service) Close() {
|
||||
// Add cleanup code if needed
|
||||
}
|
||||
245
internal/chatlog/http/route.go
Normal file
245
internal/chatlog/http/route.go
Normal file
@@ -0,0 +1,245 @@
|
||||
package http
|
||||
|
||||
import (
|
||||
"embed"
|
||||
"fmt"
|
||||
"io/fs"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/sjzar/chatlog/internal/errors"
|
||||
"github.com/sjzar/chatlog/pkg/util"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// EFS holds embedded file system data for static assets.
|
||||
//
|
||||
//go:embed static
|
||||
var EFS embed.FS
|
||||
|
||||
// initRouter sets up routes and static file servers for the web service.
|
||||
// It defines endpoints for API as well as serving static content.
|
||||
func (s *Service) initRouter() {
|
||||
|
||||
router := s.GetRouter()
|
||||
|
||||
staticDir, _ := fs.Sub(EFS, "static")
|
||||
router.StaticFS("/static", http.FS(staticDir))
|
||||
router.StaticFileFS("/favicon.ico", "./favicon.ico", http.FS(staticDir))
|
||||
router.StaticFileFS("/", "./index.htm", http.FS(staticDir))
|
||||
|
||||
// MCP Server
|
||||
{
|
||||
router.GET("/sse", s.mcp.HandleSSE)
|
||||
router.POST("/messages", s.mcp.HandleMessages)
|
||||
// mcp inspector is shit
|
||||
// https://github.com/modelcontextprotocol/inspector/blob/aeaf32f/server/src/index.ts#L155
|
||||
router.POST("/message", s.mcp.HandleMessages)
|
||||
}
|
||||
|
||||
// API V1 Router
|
||||
api := router.Group("/api/v1")
|
||||
{
|
||||
api.GET("/chatlog", s.GetChatlog)
|
||||
api.GET("/contact", s.ListContact)
|
||||
api.GET("/chatroom", s.ListChatRoom)
|
||||
api.GET("/session", s.GetSession)
|
||||
}
|
||||
|
||||
router.NoRoute(s.NoRoute)
|
||||
}
|
||||
|
||||
// NoRoute handles 404 Not Found errors. If the request URL starts with "/api"
|
||||
// or "/static", it responds with a JSON error. Otherwise, it redirects to the root path.
|
||||
func (s *Service) NoRoute(c *gin.Context) {
|
||||
path := c.Request.URL.Path
|
||||
switch {
|
||||
case strings.HasPrefix(path, "/api"), strings.HasPrefix(path, "/static"):
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "Not found"})
|
||||
default:
|
||||
c.Header("Cache-Control", "no-cache, no-store, max-age=0, must-revalidate, value")
|
||||
c.Redirect(http.StatusFound, "/")
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) GetChatlog(c *gin.Context) {
|
||||
|
||||
var err error
|
||||
start, end, ok := util.TimeRangeOf(c.Query("time"))
|
||||
if !ok {
|
||||
errors.Err(c, errors.ErrInvalidArg("time"))
|
||||
}
|
||||
|
||||
var limit int
|
||||
if _limit := c.Query("limit"); len(_limit) > 0 {
|
||||
limit, err = strconv.Atoi(_limit)
|
||||
if err != nil {
|
||||
errors.Err(c, errors.ErrInvalidArg("limit"))
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
var offset int
|
||||
if _offset := c.Query("offset"); len(_offset) > 0 {
|
||||
offset, err = strconv.Atoi(_offset)
|
||||
if err != nil {
|
||||
errors.Err(c, errors.ErrInvalidArg("offset"))
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
talker := c.Query("talker")
|
||||
|
||||
if limit < 0 {
|
||||
limit = 0
|
||||
}
|
||||
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
|
||||
messages, err := s.db.GetMessages(start, end, talker, limit, offset)
|
||||
if err != nil {
|
||||
errors.Err(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
switch strings.ToLower(c.Query("format")) {
|
||||
case "csv":
|
||||
case "json":
|
||||
// json
|
||||
c.JSON(http.StatusOK, messages)
|
||||
default:
|
||||
// plain text
|
||||
c.Writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
|
||||
c.Writer.Header().Set("Cache-Control", "no-cache")
|
||||
c.Writer.Header().Set("Connection", "keep-alive")
|
||||
c.Writer.Flush()
|
||||
|
||||
for _, m := range messages {
|
||||
c.Writer.WriteString(m.PlainText(len(talker) == 0))
|
||||
c.Writer.WriteString("\n")
|
||||
c.Writer.Flush()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) ListContact(c *gin.Context) {
|
||||
list, err := s.db.ListContact()
|
||||
if err != nil {
|
||||
errors.Err(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
format := strings.ToLower(c.Query("format"))
|
||||
switch format {
|
||||
case "json":
|
||||
// json
|
||||
c.JSON(http.StatusOK, list)
|
||||
default:
|
||||
// csv
|
||||
if format == "csv" {
|
||||
// 浏览器访问时，会下载文件
|
||||
c.Writer.Header().Set("Content-Type", "text/csv; charset=utf-8")
|
||||
} else {
|
||||
c.Writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
|
||||
}
|
||||
c.Writer.Header().Set("Cache-Control", "no-cache")
|
||||
c.Writer.Header().Set("Connection", "keep-alive")
|
||||
c.Writer.Flush()
|
||||
|
||||
c.Writer.WriteString("UserName,Alias,Remark,NickName\n")
|
||||
for _, contact := range list.Items {
|
||||
c.Writer.WriteString(fmt.Sprintf("%s,%s,%s,%s\n", contact.UserName, contact.Alias, contact.Remark, contact.NickName))
|
||||
}
|
||||
c.Writer.Flush()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) ListChatRoom(c *gin.Context) {
|
||||
|
||||
if query := c.Query("query"); len(query) > 0 {
|
||||
chatRoom := s.db.GetChatRoom(query)
|
||||
if chatRoom != nil {
|
||||
c.JSON(http.StatusOK, chatRoom)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
list, err := s.db.ListChatRoom()
|
||||
if err != nil {
|
||||
errors.Err(c, err)
|
||||
return
|
||||
}
|
||||
format := strings.ToLower(c.Query("format"))
|
||||
switch format {
|
||||
case "json":
|
||||
// json
|
||||
c.JSON(http.StatusOK, list)
|
||||
default:
|
||||
// csv
|
||||
if format == "csv" {
|
||||
// 浏览器访问时，会下载文件
|
||||
c.Writer.Header().Set("Content-Type", "text/csv; charset=utf-8")
|
||||
} else {
|
||||
c.Writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
|
||||
}
|
||||
c.Writer.Header().Set("Cache-Control", "no-cache")
|
||||
c.Writer.Header().Set("Connection", "keep-alive")
|
||||
c.Writer.Flush()
|
||||
|
||||
c.Writer.WriteString("Name,Owner,UserCount\n")
|
||||
for _, chatRoom := range list.Items {
|
||||
c.Writer.WriteString(fmt.Sprintf("%s,%s,%d\n", chatRoom.Name, chatRoom.Owner, len(chatRoom.Users)))
|
||||
}
|
||||
c.Writer.Flush()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) GetSession(c *gin.Context) {
|
||||
|
||||
var err error
|
||||
var limit int
|
||||
if _limit := c.Query("limit"); len(_limit) > 0 {
|
||||
limit, err = strconv.Atoi(_limit)
|
||||
if err != nil {
|
||||
errors.Err(c, errors.ErrInvalidArg("limit"))
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
sessions, err := s.db.GetSession(limit)
|
||||
if err != nil {
|
||||
errors.Err(c, err)
|
||||
return
|
||||
}
|
||||
format := strings.ToLower(c.Query("format"))
|
||||
switch format {
|
||||
case "csv":
|
||||
c.Writer.Header().Set("Content-Type", "text/csv; charset=utf-8")
|
||||
c.Writer.Header().Set("Cache-Control", "no-cache")
|
||||
c.Writer.Header().Set("Connection", "keep-alive")
|
||||
c.Writer.Flush()
|
||||
|
||||
c.Writer.WriteString("UserName,NOrder,NickName,Content,NTime\n")
|
||||
for _, session := range sessions.Items {
|
||||
c.Writer.WriteString(fmt.Sprintf("%s,%d,%s,%s,%s\n", session.UserName, session.NOrder, session.NickName, strings.ReplaceAll(session.Content, "\n", "\\n"), session.NTime))
|
||||
}
|
||||
c.Writer.Flush()
|
||||
case "json":
|
||||
// json
|
||||
c.JSON(http.StatusOK, sessions)
|
||||
default:
|
||||
c.Writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
|
||||
c.Writer.Header().Set("Cache-Control", "no-cache")
|
||||
c.Writer.Header().Set("Connection", "keep-alive")
|
||||
c.Writer.Flush()
|
||||
for _, session := range sessions.Items {
|
||||
c.Writer.WriteString(session.PlainText(120))
|
||||
c.Writer.WriteString("\n")
|
||||
}
|
||||
c.Writer.Flush()
|
||||
}
|
||||
}
|
||||
90
internal/chatlog/http/service.go
Normal file
90
internal/chatlog/http/service.go
Normal file
@@ -0,0 +1,90 @@
|
||||
package http
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/sjzar/chatlog/internal/chatlog/ctx"
|
||||
"github.com/sjzar/chatlog/internal/chatlog/database"
|
||||
"github.com/sjzar/chatlog/internal/chatlog/mcp"
|
||||
"github.com/sjzar/chatlog/internal/errors"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
type Service struct {
|
||||
ctx *ctx.Context
|
||||
db *database.Service
|
||||
mcp *mcp.Service
|
||||
|
||||
router *gin.Engine
|
||||
server *http.Server
|
||||
}
|
||||
|
||||
func NewService(ctx *ctx.Context, db *database.Service, mcp *mcp.Service) *Service {
|
||||
gin.SetMode(gin.ReleaseMode)
|
||||
router := gin.New()
|
||||
|
||||
// Handle error from SetTrustedProxies
|
||||
if err := router.SetTrustedProxies(nil); err != nil {
|
||||
log.Error("Failed to set trusted proxies:", err)
|
||||
}
|
||||
|
||||
// Middleware
|
||||
router.Use(
|
||||
gin.Recovery(),
|
||||
gin.LoggerWithWriter(log.StandardLogger().Out),
|
||||
)
|
||||
|
||||
s := &Service{
|
||||
ctx: ctx,
|
||||
db: db,
|
||||
mcp: mcp,
|
||||
router: router,
|
||||
}
|
||||
|
||||
s.initRouter()
|
||||
return s
|
||||
}
|
||||
|
||||
func (s *Service) Start() error {
|
||||
s.server = &http.Server{
|
||||
Addr: s.ctx.HTTPAddr,
|
||||
Handler: s.router,
|
||||
}
|
||||
|
||||
go func() {
|
||||
// Handle error from Run
|
||||
if err := s.server.ListenAndServe(); err != nil {
|
||||
log.Error("Server Stopped: ", err)
|
||||
}
|
||||
}()
|
||||
|
||||
log.Info("Server started on ", s.ctx.HTTPAddr)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Service) Stop() error {
|
||||
|
||||
if s.server == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
// 使用超时上下文优雅关闭
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
if err := s.server.Shutdown(ctx); err != nil {
|
||||
return errors.HTTP("HTTP server shutdown error", err)
|
||||
}
|
||||
|
||||
log.Info("HTTP server stopped")
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Service) GetRouter() *gin.Engine {
|
||||
return s.router
|
||||
}
|
||||
156
internal/chatlog/http/static/index.htm
Normal file
156
internal/chatlog/http/static/index.htm
Normal file
@@ -0,0 +1,156 @@
|
||||
|
||||
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>Chatlog</title>
|
||||
<style>
|
||||
.random-paragraph {
|
||||
display: none;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div id="paragraphContainer">
|
||||
<pre style="float:left;white-space:pre-wrap;" class="random-paragraph">
|
||||
('-. .-. ('-. .-') _
|
||||
( OO ) / ( OO ).-. ( OO) )
|
||||
.-----. ,--. ,--. / . --. / / '._ ,--. .-'),-----. ,----.
|
||||
' .--./ | | | | | \-. \ |'--...__) | |.-') ( OO' .-. ' ' .-./-')
|
||||
| |('-. | .| | .-'-' | | '--. .--' | | OO ) / | | | | | |_( O- )
|
||||
/_) |OO ) | | \| |_.' | | | | |`-' | \_) | |\| | | | .--, \
|
||||
|| |`-'| | .-. | | .-. | | | (| '---.' \ | | | |(| | '. (_/
|
||||
(_' '--'\ | | | | | | | | | | | | `' '-' ' | '--' |
|
||||
`-----' `--' `--' `--' `--' `--' `------' `-----' `------'
|
||||
</pre>
|
||||
<pre style="float:left;white-space:pre-wrap;" class="random-paragraph">
|
||||
_____ _ _ _
|
||||
/ __ \| | | | | |
|
||||
| / \/| |__ __ _ | |_ | | ___ __ _
|
||||
| | | '_ \ / _` || __|| | / _ \ / _` |
|
||||
| \__/\| | | || (_| || |_ | || (_) || (_| |
|
||||
\____/|_| |_| \__,_| \__||_| \___/ \__, |
|
||||
__/ |
|
||||
|___/
|
||||
</pre>
|
||||
<pre style="float:left;white-space:pre-wrap;" class="random-paragraph">
|
||||
|
||||
,-----. ,--. ,--. ,--.
|
||||
' .--./ | ,---. ,--,--. ,-' '-. | | ,---. ,---.
|
||||
| | | .-. | ' ,-. | '-. .-' | | | .-. | | .-. |
|
||||
' '--'\ | | | | \ '-' | | | | | ' '-' ' ' '-' '
|
||||
`-----' `--' `--' `--`--' `--' `--' `---' .`- /
|
||||
`---'
|
||||
</pre>
|
||||
<pre style="float:left;white-space:pre-wrap;" class="random-paragraph">
|
||||
____ _ _ _
|
||||
/ ___| | |__ __ _ | |_ | | ___ __ _
|
||||
| | | '_ \ / _` | | __| | | / _ \ / _` |
|
||||
| |___ | | | | | (_| | | |_ | | | (_) | | (_| |
|
||||
\____| |_| |_| \__,_| \__| |_| \___/ \__, |
|
||||
|___/
|
||||
</pre>
|
||||
<pre style="float:left;white-space:pre-wrap;" class="random-paragraph">
|
||||
_____ _____ _____ _____ _____ _______ _____
|
||||
/\ \ /\ \ /\ \ /\ \ /\ \ /::\ \ /\ \
|
||||
/::\ \ /::\____\ /::\ \ /::\ \ /::\____\ /::::\ \ /::\ \
|
||||
/::::\ \ /:::/ / /::::\ \ \:::\ \ /:::/ / /::::::\ \ /::::\ \
|
||||
/::::::\ \ /:::/ / /::::::\ \ \:::\ \ /:::/ / /::::::::\ \ /::::::\ \
|
||||
/:::/\:::\ \ /:::/ / /:::/\:::\ \ \:::\ \ /:::/ / /:::/~~\:::\ \ /:::/\:::\ \
|
||||
/:::/ \:::\ \ /:::/____/ /:::/__\:::\ \ \:::\ \ /:::/ / /:::/ \:::\ \ /:::/ \:::\ \
|
||||
/:::/ \:::\ \ /::::\ \ /::::\ \:::\ \ /::::\ \ /:::/ / /:::/ / \:::\ \ /:::/ \:::\ \
|
||||
/:::/ / \:::\ \ /::::::\ \ _____ /::::::\ \:::\ \ /::::::\ \ /:::/ / /:::/____/ \:::\____\ /:::/ / \:::\ \
|
||||
/:::/ / \:::\ \ /:::/\:::\ \ /\ \ /:::/\:::\ \:::\ \ /:::/\:::\ \ /:::/ / |:::| | |:::| | /:::/ / \:::\ ___\
|
||||
/:::/____/ \:::\____\/:::/ \:::\ /::\____\/:::/ \:::\ \:::\____\ /:::/ \:::\____\/:::/____/ |:::|____| |:::| |/:::/____/ ___\:::| |
|
||||
\:::\ \ \::/ /\::/ \:::\ /:::/ /\::/ \:::\ /:::/ / /:::/ \::/ /\:::\ \ \:::\ \ /:::/ / \:::\ \ /\ /:::|____|
|
||||
\:::\ \ \/____/ \/____/ \:::\/:::/ / \/____/ \:::\/:::/ / /:::/ / \/____/ \:::\ \ \:::\ \ /:::/ / \:::\ /::\ \::/ /
|
||||
\:::\ \ \::::::/ / \::::::/ / /:::/ / \:::\ \ \:::\ /:::/ / \:::\ \:::\ \/____/
|
||||
\:::\ \ \::::/ / \::::/ / /:::/ / \:::\ \ \:::\__/:::/ / \:::\ \:::\____\
|
||||
\:::\ \ /:::/ / /:::/ / \::/ / \:::\ \ \::::::::/ / \:::\ /:::/ /
|
||||
\:::\ \ /:::/ / /:::/ / \/____/ \:::\ \ \::::::/ / \:::\/:::/ /
|
||||
\:::\ \ /:::/ / /:::/ / \:::\ \ \::::/ / \::::::/ /
|
||||
\:::\____\ /:::/ / /:::/ / \:::\____\ \::/____/ \::::/ /
|
||||
\::/ / \::/ / \::/ / \::/ / ~~ \::/____/
|
||||
\/____/ \/____/ \/____/ \/____/
|
||||
|
||||
</pre>
|
||||
<pre style="float:left;white-space:pre-wrap;" class="random-paragraph">
|
||||
___ _ _ __ ____ __ _____ ___
|
||||
/ __)( )_( ) /__\ (_ _)( ) ( _ ) / __)
|
||||
( (__ ) _ ( /(__)\ )( )(__ )(_)( ( (_-.
|
||||
\___)(_) (_)(__)(__) (__) (____)(_____) \___/
|
||||
</pre>
|
||||
<pre style="float:left;white-space:pre-wrap;" class="random-paragraph">
|
||||
________ ___ ___ ________ _________ ___ ________ ________
|
||||
|\ ____\ |\ \|\ \ |\ __ \ |\___ ___\ |\ \ |\ __ \ |\ ____\
|
||||
\ \ \___| \ \ \\\ \ \ \ \|\ \ \|___ \ \_| \ \ \ \ \ \|\ \ \ \ \___|
|
||||
\ \ \ \ \ __ \ \ \ __ \ \ \ \ \ \ \ \ \ \\\ \ \ \ \ ___
|
||||
\ \ \____ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \____ \ \ \\\ \ \ \ \|\ \
|
||||
\ \_______\ \ \__\ \__\ \ \__\ \__\ \ \__\ \ \_______\ \ \_______\ \ \_______\
|
||||
\|_______| \|__|\|__| \|__|\|__| \|__| \|_______| \|_______| \|_______|
|
||||
</pre>
|
||||
<pre style="float:left;white-space:pre-wrap;" class="random-paragraph">
|
||||
╔═╗┬ ┬┌─┐┌┬┐┬ ┌─┐┌─┐
|
||||
║ ├─┤├─┤ │ │ │ ││ ┬
|
||||
╚═╝┴ ┴┴ ┴ ┴ ┴─┘└─┘└─┘
|
||||
</pre>
|
||||
<pre style="float:left;white-space:pre-wrap;" class="random-paragraph">
|
||||
▄▄· ▄ .▄ ▄▄▄· ▄▄▄▄▄▄▄▌ ▄▄ •
|
||||
▐█ ▌▪██▪▐█▐█ ▀█ •██ ██• ▪ ▐█ ▀ ▪
|
||||
██ ▄▄██▀▐█▄█▀▀█ ▐█.▪██▪ ▄█▀▄ ▄█ ▀█▄
|
||||
▐███▌██▌▐▀▐█ ▪▐▌ ▐█▌·▐█▌▐▌▐█▌.▐▌▐█▄▪▐█
|
||||
·▀▀▀ ▀▀▀ · ▀ ▀ ▀▀▀ .▀▀▀ ▀█▄▀▪·▀▀▀▀
|
||||
</pre>
|
||||
<pre style="float:left;white-space:pre-wrap;" class="random-paragraph">
|
||||
,. - ., .·¨'`; ,.·´¨;\ ,., ' , . ., ° ,. ' , ·. ,.-·~·., ‘ ,.-·^*ª'` ·,
|
||||
,·'´ ,. - , ';\ '; ;'\ '; ;::\ ;´ '· ., ;'´ , ., _';\' / ';\ / ·'´,.-·-., `,'‚ .·´ ,·'´:¯'`·, '\‘
|
||||
,·´ .'´\:::::;' ;:'\ ' ; ;::'\ ,' ;::'; .´ .-, ';\ \:´¨¯:;' `;::'\:'\ ,' ,'::'\ / .'´\:::::::'\ '\ ° ,´ ,'\:::::::::\,.·\'
|
||||
/ ,'´::::'\;:-/ ,' ::; ' ; ;::_';,. ,.' ;:::';° / /:\:'; ;:'\' \::::; ,'::_'\;' ,' ;:::';' ,·' ,'::::\:;:-·-:'; ';\‚ / /:::\;·'´¯'`·;\:::\°
|
||||
,' ;':::::;'´ '; /\::;' ' .' ,. -·~-·, ;:::'; ' ,' ,'::::'\'; ;::'; ,' ,'::;' ‘ '; ,':::;' ;. ';:::;´ ,' ,':'\‚ ; ;:::;' '\;:·´
|
||||
; ;:::::; '\*'´\::\' ° '; ;'\::::::::; '/::::; ,.-·' '·~^*'´¨, ';::; ; ;:::; ° ; ,':::;' ' '; ;::; ,'´ .'´\::';‚ '; ;::/ ,·´¯'; °
|
||||
'; ';::::'; '\::'\/.' ; ';:;\;::-··; ;::::; ':, ,·:²*´¨¯'`; ;::'; ; ;::;' ‘ ,' ,'::;' '; ':;: ,.·´,.·´::::\;'° '; '·;' ,.·´, ;'\
|
||||
\ '·:;:'_ ,. -·'´.·´\‘ ':,.·´\;' ;' ,' :::/ ' ,' / \::::::::'; ;::'; ; ;::;'‚ ; ';_:,.-·´';\‘ \·, `*´,.·'´::::::;·´ \'·. `'´,.·:´'; ;::\'
|
||||
'\:` · .,. -·:´::::::\' \:::::\ \·.'::::; ,' ,'::::\·²*'´¨¯':,'\:; ',.'\::;'‚ ', _,.-·'´:\:\‘ \\:¯::\:::::::;:·´ '\::\¯::::::::'; ;::'; ‘
|
||||
\:::::::\:::::::;:·'´' \;:·´ \:\::'; \`¨\:::/ \::\' \::\:;'‚ \¨:::::::::::\'; `\:::::\;::·'´ ° `·:\:::;:·´';.·´\::;'
|
||||
`· :;::\;::-·´ `·\;' '\::\;' '\;' ' \;:' ‘ '\;::_;:-·'´‘ ¯ ¯ \::::\;'‚
|
||||
' `¨' ° '¨ ‘ '\:·´'
|
||||
</pre>
|
||||
<pre style="float:left;white-space:pre-wrap;" class="random-paragraph">
|
||||
▄████▄ ██░ ██ ▄▄▄ ▄▄▄█████▓ ██▓ ▒█████ ▄████
|
||||
▒██▀ ▀█ ▓██░ ██▒▒████▄ ▓ ██▒ ▓▒▓██▒ ▒██▒ ██▒ ██▒ ▀█▒
|
||||
▒▓█ ▄ ▒██▀▀██░▒██ ▀█▄ ▒ ▓██░ ▒░▒██░ ▒██░ ██▒▒██░▄▄▄░
|
||||
▒▓▓▄ ▄██▒░▓█ ░██ ░██▄▄▄▄██ ░ ▓██▓ ░ ▒██░ ▒██ ██░░▓█ ██▓
|
||||
▒ ▓███▀ ░░▓█▒░██▓ ▓█ ▓██▒ ▒██▒ ░ ░██████▒░ ████▓▒░░▒▓███▀▒
|
||||
░ ░▒ ▒ ░ ▒ ░░▒░▒ ▒▒ ▓▒█░ ▒ ░░ ░ ▒░▓ ░░ ▒░▒░▒░ ░▒ ▒
|
||||
░ ▒ ▒ ░▒░ ░ ▒ ▒▒ ░ ░ ░ ░ ▒ ░ ░ ▒ ▒░ ░ ░
|
||||
░ ░ ░░ ░ ░ ▒ ░ ░ ░ ░ ░ ░ ▒ ░ ░ ░
|
||||
░ ░ ░ ░ ░ ░ ░ ░ ░ ░ ░ ░
|
||||
░
|
||||
</pre>
|
||||
<pre style="float:left;white-space:pre-wrap;" class="random-paragraph">
|
||||
▄█▄ ▄ █ ██ ▄▄▄▄▀ █ ████▄ ▄▀
|
||||
█▀ ▀▄ █ █ █ █ ▀▀▀ █ █ █ █ ▄▀
|
||||
█ ▀ ██▀▀█ █▄▄█ █ █ █ █ █ ▀▄
|
||||
█▄ ▄▀ █ █ █ █ █ ███▄ ▀████ █ █
|
||||
▀███▀ █ █ ▀ ▀ ███
|
||||
▀ █
|
||||
▀
|
||||
</pre>
|
||||
</div>
|
||||
<script>
|
||||
window.onload = function() {
|
||||
showRandomParagraph();
|
||||
};
|
||||
function showRandomParagraph() {
|
||||
const paragraphs = document.getElementsByClassName("random-paragraph");
|
||||
for (let i = 0; i < paragraphs.length; i++) {
|
||||
paragraphs[i].style.display = "none";
|
||||
}
|
||||
const randomIndex = Math.floor(Math.random() * paragraphs.length);
|
||||
paragraphs[randomIndex].style.display = "block";
|
||||
}
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
202
internal/chatlog/manager.go
Normal file
202
internal/chatlog/manager.go
Normal file
@@ -0,0 +1,202 @@
|
||||
package chatlog
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
|
||||
"github.com/sjzar/chatlog/internal/chatlog/conf"
|
||||
"github.com/sjzar/chatlog/internal/chatlog/ctx"
|
||||
"github.com/sjzar/chatlog/internal/chatlog/database"
|
||||
"github.com/sjzar/chatlog/internal/chatlog/http"
|
||||
"github.com/sjzar/chatlog/internal/chatlog/mcp"
|
||||
"github.com/sjzar/chatlog/internal/chatlog/wechat"
|
||||
"github.com/sjzar/chatlog/pkg/util"
|
||||
)
|
||||
|
||||
// Manager 管理聊天日志应用
|
||||
type Manager struct {
|
||||
conf *conf.Service
|
||||
ctx *ctx.Context
|
||||
|
||||
// Services
|
||||
db *database.Service
|
||||
http *http.Service
|
||||
mcp *mcp.Service
|
||||
wechat *wechat.Service
|
||||
|
||||
// Terminal UI
|
||||
app *App
|
||||
}
|
||||
|
||||
func New(configPath string) (*Manager, error) {
|
||||
|
||||
// 创建配置服务
|
||||
conf, err := conf.NewService(configPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 创建应用上下文
|
||||
ctx := ctx.New(conf)
|
||||
|
||||
wechat := wechat.NewService(ctx)
|
||||
|
||||
db := database.NewService(ctx)
|
||||
|
||||
mcp := mcp.NewService(ctx, db)
|
||||
|
||||
http := http.NewService(ctx, db, mcp)
|
||||
|
||||
return &Manager{
|
||||
conf: conf,
|
||||
ctx: ctx,
|
||||
db: db,
|
||||
mcp: mcp,
|
||||
http: http,
|
||||
wechat: wechat,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (m *Manager) Run() error {
|
||||
|
||||
m.ctx.WeChatInstances = m.wechat.GetWeChatInstances()
|
||||
if len(m.ctx.WeChatInstances) >= 1 {
|
||||
m.ctx.SwitchCurrent(m.ctx.WeChatInstances[0])
|
||||
}
|
||||
|
||||
if m.ctx.HTTPEnabled {
|
||||
// 启动HTTP服务
|
||||
if err := m.StartService(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
// 启动终端UI
|
||||
m.app = NewApp(m.ctx, m)
|
||||
m.app.Run() // 阻塞
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *Manager) StartService() error {
|
||||
|
||||
// 按依赖顺序启动服务
|
||||
if err := m.db.Start(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if err := m.mcp.Start(); err != nil {
|
||||
m.db.Stop() // 回滚已启动的服务
|
||||
return err
|
||||
}
|
||||
|
||||
if err := m.http.Start(); err != nil {
|
||||
m.mcp.Stop() // 回滚已启动的服务
|
||||
m.db.Stop()
|
||||
return err
|
||||
}
|
||||
|
||||
// 更新状态
|
||||
m.ctx.SetHTTPEnabled(true)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *Manager) StopService() error {
|
||||
// 按依赖的反序停止服务
|
||||
var errs []error
|
||||
|
||||
if err := m.http.Stop(); err != nil {
|
||||
errs = append(errs, err)
|
||||
}
|
||||
|
||||
if err := m.mcp.Stop(); err != nil {
|
||||
errs = append(errs, err)
|
||||
}
|
||||
|
||||
if err := m.db.Stop(); err != nil {
|
||||
errs = append(errs, err)
|
||||
}
|
||||
|
||||
// 更新状态
|
||||
m.ctx.SetHTTPEnabled(false)
|
||||
|
||||
// 如果有错误，返回第一个错误
|
||||
if len(errs) > 0 {
|
||||
return errs[0]
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *Manager) GetDataKey() error {
|
||||
if m.ctx.Current == nil {
|
||||
return fmt.Errorf("未选择任何账号")
|
||||
}
|
||||
if _, err := m.wechat.GetDataKey(m.ctx.Current); err != nil {
|
||||
return err
|
||||
}
|
||||
m.ctx.Refresh()
|
||||
m.ctx.UpdateConfig()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *Manager) DecryptDBFiles() error {
|
||||
if m.ctx.DataKey == "" {
|
||||
if m.ctx.Current == nil {
|
||||
return fmt.Errorf("未选择任何账号")
|
||||
}
|
||||
if err := m.GetDataKey(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if m.ctx.WorkDir == "" {
|
||||
m.ctx.WorkDir = util.DefaultWorkDir(m.ctx.Account)
|
||||
}
|
||||
|
||||
if err := m.wechat.DecryptDBFiles(m.ctx.DataDir, m.ctx.WorkDir, m.ctx.DataKey, m.ctx.MajorVersion); err != nil {
|
||||
return err
|
||||
}
|
||||
m.ctx.Refresh()
|
||||
m.ctx.UpdateConfig()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *Manager) CommandKey(pid int) (string, error) {
|
||||
instances := m.wechat.GetWeChatInstances()
|
||||
if len(instances) == 0 {
|
||||
return "", fmt.Errorf("wechat process not found")
|
||||
}
|
||||
if len(instances) == 1 {
|
||||
return instances[0].GetKey()
|
||||
}
|
||||
if pid == 0 {
|
||||
str := "Select a process:\n"
|
||||
for _, ins := range instances {
|
||||
str += fmt.Sprintf("PID: %d. %s[Version: %s Data Dir: %s ]\n", ins.PID, ins.AccountName, ins.Version.FileVersion, ins.DataDir)
|
||||
}
|
||||
return str, nil
|
||||
}
|
||||
for _, ins := range instances {
|
||||
if ins.PID == uint32(pid) {
|
||||
return ins.GetKey()
|
||||
}
|
||||
}
|
||||
return "", fmt.Errorf("wechat process not found")
|
||||
}
|
||||
|
||||
func (m *Manager) CommandDecrypt(dataDir string, workDir string, key string, version int) error {
|
||||
if dataDir == "" {
|
||||
return fmt.Errorf("dataDir is required")
|
||||
}
|
||||
if key == "" {
|
||||
return fmt.Errorf("key is required")
|
||||
}
|
||||
if workDir == "" {
|
||||
workDir = util.DefaultWorkDir(filepath.Base(filepath.Dir(dataDir)))
|
||||
}
|
||||
|
||||
if err := m.wechat.DecryptDBFiles(dataDir, workDir, key, version); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
96
internal/chatlog/mcp/const.go
Normal file
96
internal/chatlog/mcp/const.go
Normal file
@@ -0,0 +1,96 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"github.com/sjzar/chatlog/internal/mcp"
|
||||
)
|
||||
|
||||
// MCPTools 和资源定义
|
||||
var (
|
||||
InitializeResponse = mcp.InitializeResponse{
|
||||
ProtocolVersion: mcp.ProtocolVersion,
|
||||
Capabilities: mcp.DefaultCapabilities,
|
||||
ServerInfo: mcp.ServerInfo{
|
||||
Name: "chatlog",
|
||||
Version: "0.0.1",
|
||||
},
|
||||
}
|
||||
|
||||
ToolContact = mcp.Tool{
|
||||
Name: "query_contact",
|
||||
Description: "查询用户的联系人信息。可以通过姓名、备注名或ID进行查询，返回匹配的联系人列表。当用户询问某人的联系方式、想了解联系人信息或需要查找特定联系人时使用此工具。参数为空时，将返回联系人列表",
|
||||
InputSchema: mcp.ToolSchema{
|
||||
Type: "object",
|
||||
Properties: mcp.M{
|
||||
"query": mcp.M{
|
||||
"type": "string",
|
||||
"description": "联系人的搜索关键词，可以是姓名、备注名或ID。",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
ToolChatRoom = mcp.Tool{
|
||||
Name: "query_chat_room",
|
||||
Description: "查询用户参与的群聊信息。可以通过群名称、群ID或相关关键词进行查询，返回匹配的群聊列表。当用户询问群聊信息、想了解某个群的详情或需要查找特定群聊时使用此工具。",
|
||||
InputSchema: mcp.ToolSchema{
|
||||
Type: "object",
|
||||
Properties: mcp.M{
|
||||
"query": mcp.M{
|
||||
"type": "string",
|
||||
"description": "群聊的搜索关键词，可以是群名称、群ID或相关描述",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
ToolRecentChat = mcp.Tool{
|
||||
Name: "query_recent_chat",
|
||||
Description: "查询最近会话列表，包括个人聊天和群聊。当用户想了解最近的聊天记录、查看最近联系过的人或群组时使用此工具。不需要参数，直接返回最近的会话列表。",
|
||||
InputSchema: mcp.ToolSchema{
|
||||
Type: "object",
|
||||
Properties: mcp.M{},
|
||||
},
|
||||
}
|
||||
|
||||
ToolChatLog = mcp.Tool{
|
||||
Name: "chatlog",
|
||||
Description: "查询特定时间或时间段内与特定联系人或群组的聊天记录。当用户需要回顾过去的对话内容、查找特定信息或想了解与某人/某群的历史交流时使用此工具。",
|
||||
InputSchema: mcp.ToolSchema{
|
||||
Type: "object",
|
||||
Properties: mcp.M{
|
||||
"time": mcp.M{
|
||||
"type": "string",
|
||||
"description": "查询的时间点或时间段。可以是具体时间，例如 YYYY-MM-DD，也可以是时间段，例如 YYYY-MM-DD~YYYY-MM-DD，时间段之间用\"~\"分隔。",
|
||||
},
|
||||
"talker": mcp.M{
|
||||
"type": "string",
|
||||
"description": "交谈对象，可以是联系人或群聊。支持使用ID、昵称、备注名等进行查询。",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
ResourceRecentChat = mcp.Resource{
|
||||
Name: "最近会话",
|
||||
URI: "session://recent",
|
||||
Description: "获取最近的聊天会话列表",
|
||||
}
|
||||
|
||||
ResourceTemplateContact = mcp.ResourceTemplate{
|
||||
Name: "联系人信息",
|
||||
URITemplate: "contact://{username}",
|
||||
Description: "获取指定联系人的详细信息",
|
||||
}
|
||||
|
||||
ResourceTemplateChatRoom = mcp.ResourceTemplate{
|
||||
Name: "群聊信息",
|
||||
URITemplate: "chatroom://{roomid}",
|
||||
Description: "获取指定群聊的详细信息",
|
||||
}
|
||||
|
||||
ResourceTemplateChatlog = mcp.ResourceTemplate{
|
||||
Name: "聊天记录",
|
||||
URITemplate: "chatlog://{talker}/{timeframe}?limit,offset",
|
||||
Description: "获取与特定联系人或群聊的聊天记录",
|
||||
}
|
||||
)
|
||||
357
internal/chatlog/mcp/service.go
Normal file
357
internal/chatlog/mcp/service.go
Normal file
@@ -0,0 +1,357 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/sjzar/chatlog/internal/chatlog/ctx"
|
||||
"github.com/sjzar/chatlog/internal/chatlog/database"
|
||||
"github.com/sjzar/chatlog/internal/mcp"
|
||||
"github.com/sjzar/chatlog/pkg/util"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type Service struct {
|
||||
ctx *ctx.Context
|
||||
db *database.Service
|
||||
|
||||
mcp *mcp.MCP
|
||||
}
|
||||
|
||||
func NewService(ctx *ctx.Context, db *database.Service) *Service {
|
||||
return &Service{
|
||||
ctx: ctx,
|
||||
db: db,
|
||||
}
|
||||
}
|
||||
|
||||
// GetMCP 获取底层MCP实例
|
||||
func (s *Service) GetMCP() *mcp.MCP {
|
||||
return s.mcp
|
||||
}
|
||||
|
||||
// Start 启动MCP服务
|
||||
func (s *Service) Start() error {
|
||||
s.mcp = mcp.NewMCP()
|
||||
go s.worker()
|
||||
return nil
|
||||
}
|
||||
|
||||
// Stop 停止MCP服务
|
||||
func (s *Service) Stop() error {
|
||||
if s.mcp != nil {
|
||||
s.mcp.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// worker 处理MCP请求
|
||||
func (s *Service) worker() {
|
||||
for {
|
||||
select {
|
||||
case p, ok := <-s.mcp.ProcessChan:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
s.processMCP(p.Session, p.Request)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) HandleSSE(c *gin.Context) {
|
||||
s.mcp.HandleSSE(c)
|
||||
}
|
||||
|
||||
func (s *Service) HandleMessages(c *gin.Context) {
|
||||
s.mcp.HandleMessages(c)
|
||||
}
|
||||
|
||||
// processMCP 处理MCP请求
|
||||
func (s *Service) processMCP(session *mcp.Session, req *mcp.Request) {
|
||||
var err error
|
||||
switch req.Method {
|
||||
case mcp.MethodInitialize:
|
||||
err = s.initialize(session, req)
|
||||
case mcp.MethodToolsList:
|
||||
err = s.sendCustomParams(session, req, mcp.M{"tools": []mcp.Tool{
|
||||
ToolContact,
|
||||
ToolChatRoom,
|
||||
ToolRecentChat,
|
||||
ToolChatLog,
|
||||
}})
|
||||
case mcp.MethodToolsCall:
|
||||
err = s.toolsCall(session, req)
|
||||
case mcp.MethodPromptsList:
|
||||
err = s.sendCustomParams(session, req, mcp.M{"prompts": []mcp.Prompt{}})
|
||||
case mcp.MethodResourcesList:
|
||||
err = s.sendCustomParams(session, req, mcp.M{"resources": []mcp.Resource{
|
||||
ResourceRecentChat,
|
||||
}})
|
||||
case mcp.MethodResourcesTemplateList:
|
||||
err = s.sendCustomParams(session, req, mcp.M{"resourceTemplates": []mcp.ResourceTemplate{
|
||||
ResourceTemplateContact,
|
||||
ResourceTemplateChatRoom,
|
||||
ResourceTemplateChatlog,
|
||||
}})
|
||||
case mcp.MethodResourcesRead:
|
||||
err = s.resourcesRead(session, req)
|
||||
case mcp.MethodPing:
|
||||
err = s.sendCustomParams(session, req, struct{}{})
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
session.WriteError(req, err)
|
||||
}
|
||||
}
|
||||
|
||||
// initialize 处理初始化请求
|
||||
func (s *Service) initialize(session *mcp.Session, req *mcp.Request) error {
|
||||
initReq, err := parseParams[mcp.InitializeRequest](req.Params)
|
||||
if err != nil {
|
||||
return fmt.Errorf("解析初始化参数失败: %v", err)
|
||||
}
|
||||
session.SaveClientInfo(initReq.ClientInfo)
|
||||
|
||||
return session.WriteResponse(req, InitializeResponse)
|
||||
}
|
||||
|
||||
// toolsCall 处理工具调用
|
||||
func (s *Service) toolsCall(session *mcp.Session, req *mcp.Request) error {
|
||||
callReq, err := parseParams[mcp.ToolsCallRequest](req.Params)
|
||||
if err != nil {
|
||||
return fmt.Errorf("解析工具调用参数失败: %v", err)
|
||||
}
|
||||
|
||||
buf := &bytes.Buffer{}
|
||||
switch callReq.Name {
|
||||
case "query_contact":
|
||||
query := ""
|
||||
if v, ok := callReq.Arguments["query"]; ok {
|
||||
query = v.(string)
|
||||
}
|
||||
if len(query) == 0 {
|
||||
list, err := s.db.ListContact()
|
||||
if err != nil {
|
||||
return fmt.Errorf("无法获取联系人列表: %v", err)
|
||||
}
|
||||
buf.WriteString("UserName,Alias,Remark,NickName\n")
|
||||
for _, contact := range list.Items {
|
||||
buf.WriteString(fmt.Sprintf("%s,%s,%s,%s\n", contact.UserName, contact.Alias, contact.Remark, contact.NickName))
|
||||
}
|
||||
} else {
|
||||
contact := s.db.GetContact(query)
|
||||
if contact == nil {
|
||||
return fmt.Errorf("无法获取联系人: %s", query)
|
||||
}
|
||||
b, err := json.Marshal(contact)
|
||||
if err != nil {
|
||||
return fmt.Errorf("无法序列化联系人: %v", err)
|
||||
}
|
||||
buf.Write(b)
|
||||
}
|
||||
case "query_chat_room":
|
||||
query := ""
|
||||
if v, ok := callReq.Arguments["query"]; ok {
|
||||
query = v.(string)
|
||||
}
|
||||
if len(query) == 0 {
|
||||
list, err := s.db.ListChatRoom()
|
||||
if err != nil {
|
||||
return fmt.Errorf("无法获取群聊列表: %v", err)
|
||||
}
|
||||
buf.WriteString("Name,Remark,NickName,Owner,UserCount\n")
|
||||
for _, chatRoom := range list.Items {
|
||||
buf.WriteString(fmt.Sprintf("%s,%s,%s,%s,%d\n", chatRoom.Name, chatRoom.Remark, chatRoom.NickName, chatRoom.Owner, len(chatRoom.Users)))
|
||||
}
|
||||
} else {
|
||||
chatRoom := s.db.GetChatRoom(query)
|
||||
if chatRoom == nil {
|
||||
return fmt.Errorf("无法获取群聊: %s", query)
|
||||
}
|
||||
b, err := json.Marshal(chatRoom)
|
||||
if err != nil {
|
||||
return fmt.Errorf("无法序列化群聊: %v", err)
|
||||
}
|
||||
buf.Write(b)
|
||||
}
|
||||
case "query_recent_chat":
|
||||
data, err := s.db.GetSession(0)
|
||||
if err != nil {
|
||||
return fmt.Errorf("无法获取会话列表: %v", err)
|
||||
}
|
||||
for _, session := range data.Items {
|
||||
buf.WriteString(session.PlainText(120))
|
||||
buf.WriteString("\n")
|
||||
}
|
||||
case "chatlog":
|
||||
if callReq.Arguments == nil {
|
||||
return mcp.ErrInvalidParams
|
||||
}
|
||||
_time := ""
|
||||
if v, ok := callReq.Arguments["time"]; ok {
|
||||
_time = v.(string)
|
||||
}
|
||||
start, end, ok := util.TimeRangeOf(_time)
|
||||
if !ok {
|
||||
return fmt.Errorf("无法解析时间范围")
|
||||
}
|
||||
talker := ""
|
||||
if v, ok := callReq.Arguments["talker"]; ok {
|
||||
talker = v.(string)
|
||||
}
|
||||
limit := util.MustAnyToInt(callReq.Arguments["limit"])
|
||||
if limit == 0 {
|
||||
limit = 100
|
||||
}
|
||||
offset := util.MustAnyToInt(callReq.Arguments["offset"])
|
||||
messages, err := s.db.GetMessages(start, end, talker, limit, offset)
|
||||
if err != nil {
|
||||
return fmt.Errorf("无法获取聊天记录: %v", err)
|
||||
}
|
||||
for _, m := range messages {
|
||||
buf.WriteString(m.PlainText(len(talker) == 0))
|
||||
buf.WriteString("\n")
|
||||
}
|
||||
default:
|
||||
return fmt.Errorf("未支持的工具: %s", callReq.Name)
|
||||
}
|
||||
|
||||
resp := mcp.ToolsCallResponse{
|
||||
Content: []mcp.Content{
|
||||
{Type: "text", Text: buf.String()},
|
||||
},
|
||||
IsError: false,
|
||||
}
|
||||
return session.WriteResponse(req, resp)
|
||||
}
|
||||
|
||||
// resourcesRead 处理资源读取
|
||||
func (s *Service) resourcesRead(session *mcp.Session, req *mcp.Request) error {
|
||||
readReq, err := parseParams[mcp.ResourcesReadRequest](req.Params)
|
||||
if err != nil {
|
||||
return fmt.Errorf("解析资源读取参数失败: %v", err)
|
||||
}
|
||||
|
||||
u, err := url.Parse(readReq.URI)
|
||||
if err != nil {
|
||||
return fmt.Errorf("无法解析URI: %v", err)
|
||||
}
|
||||
|
||||
buf := &bytes.Buffer{}
|
||||
switch u.Scheme {
|
||||
case "contact":
|
||||
if len(u.Host) == 0 {
|
||||
list, err := s.db.ListContact()
|
||||
if err != nil {
|
||||
return fmt.Errorf("无法获取联系人列表: %v", err)
|
||||
}
|
||||
buf.WriteString("UserName,Alias,Remark,NickName\n")
|
||||
for _, contact := range list.Items {
|
||||
buf.WriteString(fmt.Sprintf("%s,%s,%s,%s\n", contact.UserName, contact.Alias, contact.Remark, contact.NickName))
|
||||
}
|
||||
} else {
|
||||
contact := s.db.GetContact(u.Host)
|
||||
if contact == nil {
|
||||
return fmt.Errorf("无法获取联系人: %s", u.Host)
|
||||
}
|
||||
b, err := json.Marshal(contact)
|
||||
if err != nil {
|
||||
return fmt.Errorf("无法序列化联系人: %v", err)
|
||||
}
|
||||
buf.Write(b)
|
||||
}
|
||||
case "chatroom":
|
||||
if len(u.Host) == 0 {
|
||||
list, err := s.db.ListChatRoom()
|
||||
if err != nil {
|
||||
return fmt.Errorf("无法获取群聊列表: %v", err)
|
||||
}
|
||||
buf.WriteString("Name,Remark,NickName,Owner,UserCount\n")
|
||||
for _, chatRoom := range list.Items {
|
||||
buf.WriteString(fmt.Sprintf("%s,%s,%s,%s,%d\n", chatRoom.Name, chatRoom.Remark, chatRoom.NickName, chatRoom.Owner, len(chatRoom.Users)))
|
||||
}
|
||||
} else {
|
||||
chatRoom := s.db.GetChatRoom(u.Host)
|
||||
if chatRoom == nil {
|
||||
return fmt.Errorf("无法获取群聊: %s", u.Host)
|
||||
}
|
||||
b, err := json.Marshal(chatRoom)
|
||||
if err != nil {
|
||||
return fmt.Errorf("无法序列化群聊: %v", err)
|
||||
}
|
||||
buf.Write(b)
|
||||
}
|
||||
case "session":
|
||||
data, err := s.db.GetSession(0)
|
||||
if err != nil {
|
||||
return fmt.Errorf("无法获取会话列表: %v", err)
|
||||
}
|
||||
for _, session := range data.Items {
|
||||
buf.WriteString(session.PlainText(120))
|
||||
buf.WriteString("\n")
|
||||
}
|
||||
case "chatlog":
|
||||
start, end, ok := util.TimeRangeOf(strings.TrimPrefix(u.Path, "/"))
|
||||
if !ok {
|
||||
return fmt.Errorf("无法解析时间范围")
|
||||
}
|
||||
limit := util.MustAnyToInt(u.Query().Get("limit"))
|
||||
if limit == 0 {
|
||||
limit = 100
|
||||
}
|
||||
offset := util.MustAnyToInt(u.Query().Get("offset"))
|
||||
messages, err := s.db.GetMessages(start, end, u.Host, limit, offset)
|
||||
if err != nil {
|
||||
return fmt.Errorf("无法获取聊天记录: %v", err)
|
||||
}
|
||||
for _, m := range messages {
|
||||
buf.WriteString(m.PlainText(len(u.Host) == 0))
|
||||
buf.WriteString("\n")
|
||||
}
|
||||
default:
|
||||
return fmt.Errorf("不支持的URI: %s", readReq.URI)
|
||||
}
|
||||
|
||||
resp := mcp.ReadingResource{
|
||||
Contents: []mcp.ReadingResourceContent{
|
||||
{URI: readReq.URI, Text: buf.String()},
|
||||
},
|
||||
}
|
||||
return session.WriteResponse(req, resp)
|
||||
}
|
||||
|
||||
// sendCustomParams 发送自定义参数
|
||||
func (s *Service) sendCustomParams(session *mcp.Session, req *mcp.Request, params interface{}) error {
|
||||
b, err := json.Marshal(mcp.NewResponse(req.ID, params))
|
||||
if err != nil {
|
||||
return fmt.Errorf("无法序列化响应: %v", err)
|
||||
}
|
||||
session.Write(b)
|
||||
return nil
|
||||
}
|
||||
|
||||
// parseParams 解析参数
|
||||
func parseParams[T any](params interface{}) (*T, error) {
|
||||
if params == nil {
|
||||
return nil, errors.New("params is nil")
|
||||
}
|
||||
|
||||
// 将 params 重新编码为 JSON
|
||||
jsonData, err := json.Marshal(params)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("无法编码 params: %v", err)
|
||||
}
|
||||
|
||||
// 解码到目标结构体
|
||||
var result T
|
||||
if err := json.Unmarshal(jsonData, &result); err != nil {
|
||||
return nil, fmt.Errorf("无法解码为目标结构体: %v", err)
|
||||
}
|
||||
|
||||
return &result, nil
|
||||
}
|
||||
113
internal/chatlog/wechat/service.go
Normal file
113
internal/chatlog/wechat/service.go
Normal file
@@ -0,0 +1,113 @@
|
||||
package wechat
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/sjzar/chatlog/internal/chatlog/ctx"
|
||||
"github.com/sjzar/chatlog/internal/wechat"
|
||||
"github.com/sjzar/chatlog/pkg/util"
|
||||
)
|
||||
|
||||
type Service struct {
|
||||
ctx *ctx.Context
|
||||
}
|
||||
|
||||
func NewService(ctx *ctx.Context) *Service {
|
||||
return &Service{
|
||||
ctx: ctx,
|
||||
}
|
||||
}
|
||||
|
||||
// GetWeChatInstances returns all running WeChat instances
|
||||
func (s *Service) GetWeChatInstances() []*wechat.Info {
|
||||
wechat.Load()
|
||||
return wechat.Items
|
||||
}
|
||||
|
||||
// GetDataKey extracts the encryption key from a WeChat process
|
||||
func (s *Service) GetDataKey(info *wechat.Info) (string, error) {
|
||||
if info == nil {
|
||||
return "", fmt.Errorf("no WeChat instance selected")
|
||||
}
|
||||
|
||||
key, err := info.GetKey()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
return key, nil
|
||||
}
|
||||
|
||||
// FindDBFiles finds all .db files in the specified directory
|
||||
func (s *Service) FindDBFiles(rootDir string, recursive bool) ([]string, error) {
|
||||
// Check if directory exists
|
||||
info, err := os.Stat(rootDir)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot access directory %s: %w", rootDir, err)
|
||||
}
|
||||
|
||||
if !info.IsDir() {
|
||||
return nil, fmt.Errorf("%s is not a directory", rootDir)
|
||||
}
|
||||
|
||||
var dbFiles []string
|
||||
|
||||
// Define walk function
|
||||
walkFunc := func(path string, info os.FileInfo, err error) error {
|
||||
if err != nil {
|
||||
// If a file or directory can't be accessed, log the error but continue
|
||||
fmt.Printf("Warning: Cannot access %s: %v\n", path, err)
|
||||
return nil
|
||||
}
|
||||
|
||||
// If it's a directory and not the root directory, and we're not recursively searching, skip it
|
||||
if info.IsDir() && path != rootDir && !recursive {
|
||||
return filepath.SkipDir
|
||||
}
|
||||
|
||||
// Check if file extension is .db
|
||||
if !info.IsDir() && strings.ToLower(filepath.Ext(path)) == ".db" {
|
||||
dbFiles = append(dbFiles, path)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Start traversal
|
||||
err = filepath.Walk(rootDir, walkFunc)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error traversing directory: %w", err)
|
||||
}
|
||||
|
||||
if len(dbFiles) == 0 {
|
||||
return nil, fmt.Errorf("no .db files found")
|
||||
}
|
||||
|
||||
return dbFiles, nil
|
||||
}
|
||||
|
||||
func (s *Service) DecryptDBFiles(dataDir string, workDir string, key string, version int) error {
|
||||
|
||||
dbfiles, err := s.FindDBFiles(dataDir, true)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for _, dbfile := range dbfiles {
|
||||
output := filepath.Join(workDir, dbfile[len(dataDir):])
|
||||
if err := util.PrepareDir(filepath.Dir(output)); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := wechat.DecryptDBFileToFile(dbfile, output, key, version); err != nil {
|
||||
if err == wechat.ErrAlreadyDecrypted {
|
||||
continue
|
||||
}
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
85
internal/errors/errors.go
Normal file
85
internal/errors/errors.go
Normal file
@@ -0,0 +1,85 @@
|
||||
package errors
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// 定义错误类型常量
|
||||
const (
|
||||
ErrTypeDatabase = "database"
|
||||
ErrTypeWeChat = "wechat"
|
||||
ErrTypeHTTP = "http"
|
||||
ErrTypeConfig = "config"
|
||||
ErrTypeInvalidArg = "invalid_argument"
|
||||
)
|
||||
|
||||
// AppError 表示应用程序错误
|
||||
type AppError struct {
|
||||
Type string `json:"type"` // 错误类型
|
||||
Message string `json:"message"` // 错误消息
|
||||
Cause error `json:"-"` // 原始错误
|
||||
Code int `json:"-"` // HTTP Code
|
||||
}
|
||||
|
||||
func (e *AppError) Error() string {
|
||||
if e.Cause != nil {
|
||||
return fmt.Sprintf("%s: %s: %v", e.Type, e.Message, e.Cause)
|
||||
}
|
||||
return fmt.Sprintf("%s: %s", e.Type, e.Message)
|
||||
}
|
||||
|
||||
func (e *AppError) String() string {
|
||||
return e.Error()
|
||||
}
|
||||
|
||||
// New 创建新的应用错误
|
||||
func New(errType, message string, cause error, code int) *AppError {
|
||||
return &AppError{
|
||||
Type: errType,
|
||||
Message: message,
|
||||
Cause: cause,
|
||||
Code: code,
|
||||
}
|
||||
}
|
||||
|
||||
// ErrInvalidArg 无效参数错误
|
||||
func ErrInvalidArg(param string) *AppError {
|
||||
return New(ErrTypeInvalidArg, fmt.Sprintf("invalid arg: %s", param), nil, http.StatusBadRequest)
|
||||
}
|
||||
|
||||
// Database 创建数据库错误
|
||||
func Database(message string, cause error) *AppError {
|
||||
return New(ErrTypeDatabase, message, cause, http.StatusInternalServerError)
|
||||
}
|
||||
|
||||
// WeChat 创建微信相关错误
|
||||
func WeChat(message string, cause error) *AppError {
|
||||
return New(ErrTypeWeChat, message, cause, http.StatusInternalServerError)
|
||||
}
|
||||
|
||||
// HTTP 创建HTTP服务错误
|
||||
func HTTP(message string, cause error) *AppError {
|
||||
return New(ErrTypeHTTP, message, cause, http.StatusInternalServerError)
|
||||
}
|
||||
|
||||
// Config 创建配置错误
|
||||
func Config(message string, cause error) *AppError {
|
||||
return New(ErrTypeConfig, message, cause, http.StatusInternalServerError)
|
||||
}
|
||||
|
||||
// Err 在HTTP响应中返回错误
|
||||
func Err(c *gin.Context, err error) {
|
||||
if appErr, ok := err.(*AppError); ok {
|
||||
c.JSON(appErr.Code, appErr)
|
||||
return
|
||||
}
|
||||
|
||||
// 未知错误
|
||||
c.JSON(http.StatusInternalServerError, gin.H{
|
||||
"type": "unknown",
|
||||
"message": err.Error(),
|
||||
})
|
||||
}
|
||||
55
internal/mcp/error.go
Normal file
55
internal/mcp/error.go
Normal file
@@ -0,0 +1,55 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
)
|
||||
|
||||
// enum ErrorCode {
|
||||
// // Standard JSON-RPC error codes
|
||||
// ParseError = -32700,
|
||||
// InvalidRequest = -32600,
|
||||
// MethodNotFound = -32601,
|
||||
// InvalidParams = -32602,
|
||||
// InternalError = -32603
|
||||
// }
|
||||
|
||||
// Error
|
||||
type Error struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
Data interface{} `json:"data,omitempty"`
|
||||
}
|
||||
|
||||
var (
|
||||
ErrParseError = &Error{Code: -32700, Message: "Parse error"}
|
||||
ErrInvalidRequest = &Error{Code: -32600, Message: "Invalid Request"}
|
||||
ErrMethodNotFound = &Error{Code: -32601, Message: "Method not found"}
|
||||
ErrInvalidParams = &Error{Code: -32602, Message: "Invalid params"}
|
||||
ErrInternalError = &Error{Code: -32603, Message: "Internal error"}
|
||||
|
||||
ErrInvalidSessionID = &Error{Code: 400, Message: "Invalid session ID"}
|
||||
ErrSessionNotFound = &Error{Code: 404, Message: "Could not find session"}
|
||||
ErrTooManyRequests = &Error{Code: 429, Message: "Too many requests"}
|
||||
)
|
||||
|
||||
func (e *Error) Error() string {
|
||||
return fmt.Sprintf("%d: %s", e.Code, e.Message)
|
||||
}
|
||||
|
||||
func (e *Error) JsonRPC() Response {
|
||||
return Response{
|
||||
JsonRPC: JsonRPCVersion,
|
||||
Error: e,
|
||||
}
|
||||
}
|
||||
|
||||
func NewErrorResponse(id interface{}, code int, err error) *Response {
|
||||
return &Response{
|
||||
JsonRPC: JsonRPCVersion,
|
||||
ID: id,
|
||||
Error: &Error{
|
||||
Code: code,
|
||||
Message: err.Error(),
|
||||
},
|
||||
}
|
||||
}
|
||||
78
internal/mcp/initialize.go
Normal file
78
internal/mcp/initialize.go
Normal file
@@ -0,0 +1,78 @@
|
||||
package mcp
|
||||
|
||||
const (
|
||||
MethodInitialize = "initialize"
|
||||
MethodPing = "ping"
|
||||
ProtocolVersion = "2024-11-05"
|
||||
)
|
||||
|
||||
// {
|
||||
// "method": "initialize",
|
||||
// "params": {
|
||||
// "protocolVersion": "2024-11-05",
|
||||
// "capabilities": {
|
||||
// "sampling": {},
|
||||
// "roots": {
|
||||
// "listChanged": true
|
||||
// }
|
||||
// },
|
||||
// "clientInfo": {
|
||||
// "name": "mcp-inspector",
|
||||
// "version": "0.0.1"
|
||||
// }
|
||||
// },
|
||||
// "jsonrpc": "2.0",
|
||||
// "id": 0
|
||||
// }
|
||||
type InitializeRequest struct {
|
||||
ProtocolVersion string `json:"protocolVersion"`
|
||||
Capabilities M `json:"capabilities"`
|
||||
ClientInfo *ClientInfo `json:"clientInfo"`
|
||||
}
|
||||
|
||||
type ClientInfo struct {
|
||||
Name string `json:"name"`
|
||||
Version string `json:"version"`
|
||||
}
|
||||
|
||||
// {
|
||||
// "jsonrpc": "2.0",
|
||||
// "id": 0,
|
||||
// "result": {
|
||||
// "protocolVersion": "2024-11-05",
|
||||
// "capabilities": {
|
||||
// "experimental": {},
|
||||
// "prompts": {
|
||||
// "listChanged": false
|
||||
// },
|
||||
// "resources": {
|
||||
// "subscribe": false,
|
||||
// "listChanged": false
|
||||
// },
|
||||
// "tools": {
|
||||
// "listChanged": false
|
||||
// }
|
||||
// },
|
||||
// "serverInfo": {
|
||||
// "name": "weather",
|
||||
// "version": "1.4.1"
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
type InitializeResponse struct {
|
||||
ProtocolVersion string `json:"protocolVersion"`
|
||||
Capabilities M `json:"capabilities"`
|
||||
ServerInfo ServerInfo `json:"serverInfo"`
|
||||
}
|
||||
|
||||
type ServerInfo struct {
|
||||
Name string `json:"name"`
|
||||
Version string `json:"version"`
|
||||
}
|
||||
|
||||
var DefaultCapabilities = M{
|
||||
"experimental": M{},
|
||||
"prompts": M{"listChanged": false},
|
||||
"resources": M{"subscribe": false, "listChanged": false},
|
||||
"tools": M{"listChanged": false},
|
||||
}
|
||||
62
internal/mcp/jsonrpc.go
Normal file
62
internal/mcp/jsonrpc.go
Normal file
@@ -0,0 +1,62 @@
|
||||
package mcp
|
||||
|
||||
const (
|
||||
JsonRPCVersion = "2.0"
|
||||
)
|
||||
|
||||
// Documents: https://modelcontextprotocol.io/docs/concepts/transports
|
||||
|
||||
// Request
|
||||
//
|
||||
// {
|
||||
// jsonrpc: "2.0",
|
||||
// id: number | string,
|
||||
// method: string,
|
||||
// params?: object
|
||||
// }
|
||||
type Request struct {
|
||||
JsonRPC string `json:"jsonrpc"`
|
||||
ID interface{} `json:"id"`
|
||||
Method string `json:"method"`
|
||||
Params interface{} `json:"params,omitempty"`
|
||||
}
|
||||
|
||||
// Response
|
||||
//
|
||||
// {
|
||||
// jsonrpc: "2.0",
|
||||
// id: number | string,
|
||||
// result?: object,
|
||||
// error?: {
|
||||
// code: number,
|
||||
// message: string,
|
||||
// data?: unknown
|
||||
// }
|
||||
// }
|
||||
type Response struct {
|
||||
JsonRPC string `json:"jsonrpc"`
|
||||
ID interface{} `json:"id"`
|
||||
Result interface{} `json:"result,omitempty"`
|
||||
Error *Error `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
func NewResponse(id interface{}, result interface{}) *Response {
|
||||
return &Response{
|
||||
JsonRPC: JsonRPCVersion,
|
||||
ID: id,
|
||||
Result: result,
|
||||
}
|
||||
}
|
||||
|
||||
// Notifications
|
||||
//
|
||||
// {
|
||||
// jsonrpc: "2.0",
|
||||
// method: string,
|
||||
// params?: object
|
||||
// }
|
||||
type Notification struct {
|
||||
JsonRPC string `json:"jsonrpc"`
|
||||
Method string `json:"method"`
|
||||
Params interface{} `json:"params,omitempty"`
|
||||
}
|
||||
107
internal/mcp/mcp.go
Normal file
107
internal/mcp/mcp.go
Normal file
@@ -0,0 +1,107 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net/http"
|
||||
"sync"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
const (
|
||||
ProcessChanCap = 1000
|
||||
)
|
||||
|
||||
type MCP struct {
|
||||
sessions map[string]*Session
|
||||
sessionMu sync.Mutex
|
||||
|
||||
ProcessChan chan ProcessCtx
|
||||
}
|
||||
|
||||
func NewMCP() *MCP {
|
||||
return &MCP{
|
||||
sessions: make(map[string]*Session),
|
||||
ProcessChan: make(chan ProcessCtx, ProcessChanCap),
|
||||
}
|
||||
}
|
||||
|
||||
func (m *MCP) HandleSSE(c *gin.Context) {
|
||||
id := uuid.New().String()
|
||||
m.sessionMu.Lock()
|
||||
m.sessions[id] = NewSession(c, id)
|
||||
m.sessionMu.Unlock()
|
||||
|
||||
c.Stream(func(w io.Writer) bool {
|
||||
<-c.Request.Context().Done()
|
||||
return false
|
||||
})
|
||||
|
||||
m.sessionMu.Lock()
|
||||
delete(m.sessions, id)
|
||||
m.sessionMu.Unlock()
|
||||
}
|
||||
|
||||
func (m *MCP) GetSession(id string) *Session {
|
||||
m.sessionMu.Lock()
|
||||
defer m.sessionMu.Unlock()
|
||||
return m.sessions[id]
|
||||
}
|
||||
|
||||
func (m *MCP) HandleMessages(c *gin.Context) {
|
||||
|
||||
// panic("xxx")
|
||||
|
||||
// 啊这, 一个 sessionid 有 3 种写法 session_id, sessionId, sessionid
|
||||
// 官方 SDK 是 session_id: https://github.com/modelcontextprotocol/python-sdk/blob/c897868/src/mcp/server/sse.py#L98
|
||||
// 写的是 sessionId: https://github.com/modelcontextprotocol/inspector/blob/aeaf32f/server/src/index.ts#L157
|
||||
|
||||
sessionID := c.Query("session_id")
|
||||
if sessionID == "" {
|
||||
sessionID = c.Query("sessionId")
|
||||
}
|
||||
if sessionID == "" {
|
||||
sessionID = c.Param("sessionid")
|
||||
}
|
||||
if sessionID == "" {
|
||||
c.JSON(http.StatusBadRequest, ErrInvalidSessionID.JsonRPC())
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
session := m.GetSession(sessionID)
|
||||
if session == nil {
|
||||
c.JSON(http.StatusNotFound, ErrSessionNotFound.JsonRPC())
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
var req Request
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, ErrInvalidRequest.JsonRPC())
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
log.Printf("收到消息: %v\n", req)
|
||||
select {
|
||||
case m.ProcessChan <- ProcessCtx{Session: session, Request: &req}:
|
||||
default:
|
||||
c.JSON(http.StatusTooManyRequests, ErrTooManyRequests.JsonRPC())
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
c.String(http.StatusAccepted, "Accepted")
|
||||
}
|
||||
|
||||
func (m *MCP) Close() {
|
||||
close(m.ProcessChan)
|
||||
}
|
||||
|
||||
type ProcessCtx struct {
|
||||
Session *Session
|
||||
Request *Request
|
||||
}
|
||||
137
internal/mcp/prompt.go
Normal file
137
internal/mcp/prompt.go
Normal file
@@ -0,0 +1,137 @@
|
||||
package mcp
|
||||
|
||||
// Document: https://modelcontextprotocol.io/docs/concepts/prompts
|
||||
|
||||
const (
|
||||
// Client => Server
|
||||
MethodPromptsList = "prompts/list"
|
||||
MethodPromptsGet = "prompts/get"
|
||||
)
|
||||
|
||||
// Prompt
|
||||
//
|
||||
// {
|
||||
// name: string; // Unique identifier for the prompt
|
||||
// description?: string; // Human-readable description
|
||||
// arguments?: [ // Optional list of arguments
|
||||
// {
|
||||
// name: string; // Argument identifier
|
||||
// description?: string; // Argument description
|
||||
// required?: boolean; // Whether argument is required
|
||||
// }
|
||||
// ]
|
||||
// }
|
||||
type Prompt struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description,omitempty"`
|
||||
Arguments []PromptArgument `json:"arguments,omitempty"`
|
||||
}
|
||||
|
||||
type PromptArgument struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description,omitempty"`
|
||||
Required bool `json:"required,omitempty"`
|
||||
}
|
||||
|
||||
// ListPrompts
|
||||
//
|
||||
// {
|
||||
// prompts: [
|
||||
// {
|
||||
// name: "analyze-code",
|
||||
// description: "Analyze code for potential improvements",
|
||||
// arguments: [
|
||||
// {
|
||||
// name: "language",
|
||||
// description: "Programming language",
|
||||
// required: true
|
||||
// }
|
||||
// ]
|
||||
// }
|
||||
// ]
|
||||
// }
|
||||
type PromptsListResponse struct {
|
||||
Prompts []Prompt `json:"prompts"`
|
||||
}
|
||||
|
||||
// Use Prompt
|
||||
// Request
|
||||
//
|
||||
// {
|
||||
// method: "prompts/get",
|
||||
// params: {
|
||||
// name: "analyze-code",
|
||||
// arguments: {
|
||||
// language: "python"
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// Response
|
||||
//
|
||||
// {
|
||||
// description: "Analyze Python code for potential improvements",
|
||||
// messages: [
|
||||
// {
|
||||
// role: "user",
|
||||
// content: {
|
||||
// type: "text",
|
||||
// text: "Please analyze the following Python code for potential improvements:\n\n```python\ndef calculate_sum(numbers):\n total = 0\n for num in numbers:\n total = total + num\n return total\n\nresult = calculate_sum([1, 2, 3, 4, 5])\nprint(result)\n```"
|
||||
// }
|
||||
// }
|
||||
// ]
|
||||
// }
|
||||
type PromptsGetRequest struct {
|
||||
Name string `json:"name"`
|
||||
Arguments M `json:"arguments"`
|
||||
}
|
||||
|
||||
type PromptsGetResponse struct {
|
||||
Description string `json:"description"`
|
||||
Messages []PromptMessage `json:"messages"`
|
||||
}
|
||||
|
||||
type PromptMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content PromptContent `json:"content"`
|
||||
}
|
||||
|
||||
type PromptContent struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text,omitempty"`
|
||||
Resource interface{} `json:"resource,omitempty"` // Resource or ResourceTemplate
|
||||
}
|
||||
|
||||
// {
|
||||
// "messages": [
|
||||
// {
|
||||
// "role": "user",
|
||||
// "content": {
|
||||
// "type": "text",
|
||||
// "text": "Analyze these system logs and the code file for any issues:"
|
||||
// }
|
||||
// },
|
||||
// {
|
||||
// "role": "user",
|
||||
// "content": {
|
||||
// "type": "resource",
|
||||
// "resource": {
|
||||
// "uri": "logs://recent?timeframe=1h",
|
||||
// "text": "[2024-03-14 15:32:11] ERROR: Connection timeout in network.py:127\n[2024-03-14 15:32:15] WARN: Retrying connection (attempt 2/3)\n[2024-03-14 15:32:20] ERROR: Max retries exceeded",
|
||||
// "mimeType": "text/plain"
|
||||
// }
|
||||
// }
|
||||
// },
|
||||
// {
|
||||
// "role": "user",
|
||||
// "content": {
|
||||
// "type": "resource",
|
||||
// "resource": {
|
||||
// "uri": "file:///path/to/code.py",
|
||||
// "text": "def connect_to_service(timeout=30):\n retries = 3\n for attempt in range(retries):\n try:\n return establish_connection(timeout)\n except TimeoutError:\n if attempt == retries - 1:\n raise\n time.sleep(5)\n\ndef establish_connection(timeout):\n # Connection implementation\n pass",
|
||||
// "mimeType": "text/x-python"
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
// ]
|
||||
// }
|
||||
74
internal/mcp/resource.go
Normal file
74
internal/mcp/resource.go
Normal file
@@ -0,0 +1,74 @@
|
||||
package mcp
|
||||
|
||||
// Document: https://modelcontextprotocol.io/docs/concepts/resources
|
||||
|
||||
const (
|
||||
// Client => Server
|
||||
MethodResourcesList = "resources/list"
|
||||
MethodResourcesTemplateList = "resources/templates/list"
|
||||
MethodResourcesRead = "resources/read"
|
||||
MethodResourcesSubscribe = "resources/subscribe"
|
||||
MethodResourcesUnsubscribe = "resources/unsubscribe"
|
||||
|
||||
// Server => Client
|
||||
NotificationResourcesListChanged = "notifications/resources/list_changed"
|
||||
NofiticationResourcesUpdated = "notifications/resources/updated"
|
||||
)
|
||||
|
||||
// Direct resources
|
||||
//
|
||||
// {
|
||||
// uri: string; // Unique identifier for the resource
|
||||
// name: string; // Human-readable name
|
||||
// description?: string; // Optional description
|
||||
// mimeType?: string; // Optional MIME type
|
||||
// }
|
||||
type Resource struct {
|
||||
URI string `json:"uri"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description,omitempty"`
|
||||
MimeType string `json:"mimeType,omitempty"`
|
||||
}
|
||||
|
||||
// Resource templates
|
||||
//
|
||||
// {
|
||||
// uriTemplate: string; // URI template following RFC 6570
|
||||
// name: string; // Human-readable name for this type
|
||||
// description?: string; // Optional description
|
||||
// mimeType?: string; // Optional MIME type for all matching resources
|
||||
// }
|
||||
type ResourceTemplate struct {
|
||||
URITemplate string `json:"uriTemplate"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description,omitempty"`
|
||||
MimeType string `json:"mimeType,omitempty"`
|
||||
}
|
||||
|
||||
// Reading resources
|
||||
// {
|
||||
// contents: [
|
||||
// {
|
||||
// uri: string; // The URI of the resource
|
||||
// mimeType?: string; // Optional MIME type
|
||||
|
||||
// // One of:
|
||||
// text?: string; // For text resources
|
||||
// blob?: string; // For binary resources (base64 encoded)
|
||||
// }
|
||||
// ]
|
||||
// }
|
||||
type ReadingResource struct {
|
||||
Contents []ReadingResourceContent `json:"contents"`
|
||||
}
|
||||
|
||||
type ResourcesReadRequest struct {
|
||||
URI string `json:"uri"`
|
||||
}
|
||||
|
||||
type ReadingResourceContent struct {
|
||||
URI string `json:"uri"`
|
||||
MimeType string `json:"mimeType,omitempty"`
|
||||
Text string `json:"text,omitempty"`
|
||||
Blob string `json:"blob,omitempty"`
|
||||
}
|
||||
48
internal/mcp/session.go
Normal file
48
internal/mcp/session.go
Normal file
@@ -0,0 +1,48 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"io"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type Session struct {
|
||||
id string
|
||||
w io.Writer
|
||||
c *ClientInfo
|
||||
}
|
||||
|
||||
func NewSession(c *gin.Context, id string) *Session {
|
||||
return &Session{
|
||||
id: id,
|
||||
w: NewSSEWriter(c, id),
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Session) Write(p []byte) (n int, err error) {
|
||||
return s.w.Write(p)
|
||||
}
|
||||
|
||||
func (s *Session) WriteError(req *Request, err error) {
|
||||
resp := NewErrorResponse(req.ID, 500, err)
|
||||
b, err := json.Marshal(resp)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
s.Write(b)
|
||||
}
|
||||
|
||||
func (s *Session) WriteResponse(req *Request, data interface{}) error {
|
||||
resp := NewResponse(req.ID, data)
|
||||
b, err := json.Marshal(resp)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s.Write(b)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Session) SaveClientInfo(c *ClientInfo) {
|
||||
s.c = c
|
||||
}
|
||||
160
internal/mcp/sse.go
Normal file
160
internal/mcp/sse.go
Normal file
@@ -0,0 +1,160 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
const (
|
||||
SSEPingIntervalS = 30
|
||||
SSEMessageChanCap = 100
|
||||
SSEContentType = "text/event-stream; charset=utf-8"
|
||||
)
|
||||
|
||||
type SSEWriter struct {
|
||||
id string
|
||||
c *gin.Context
|
||||
}
|
||||
|
||||
func NewSSEWriter(c *gin.Context, id string) *SSEWriter {
|
||||
c.Writer.Header().Set("Content-Type", SSEContentType)
|
||||
c.Writer.Header().Set("Cache-Control", "no-cache")
|
||||
c.Writer.Header().Set("Connection", "keep-alive")
|
||||
c.Writer.Header().Set("Access-Control-Allow-Origin", "*")
|
||||
c.Writer.Flush()
|
||||
|
||||
w := &SSEWriter{
|
||||
id: id,
|
||||
c: c,
|
||||
}
|
||||
w.WriteEndpoing()
|
||||
go w.ping()
|
||||
return w
|
||||
}
|
||||
|
||||
func (w *SSEWriter) Write(p []byte) (n int, err error) {
|
||||
w.WriteMessage(string(p))
|
||||
return len(p), nil
|
||||
}
|
||||
|
||||
func (w *SSEWriter) WriteMessage(data string) {
|
||||
w.WriteEvent("message", data)
|
||||
}
|
||||
|
||||
func (w *SSEWriter) WriteEvent(event string, data string) {
|
||||
w.c.Writer.WriteString(fmt.Sprintf("event: %s\n", event))
|
||||
w.c.Writer.WriteString(fmt.Sprintf("data: %s\n\n", data))
|
||||
w.c.Writer.Flush()
|
||||
}
|
||||
|
||||
func (w *SSEWriter) ping() {
|
||||
for {
|
||||
select {
|
||||
case <-time.After(time.Second * SSEPingIntervalS):
|
||||
w.writePing()
|
||||
case <-w.c.Request.Context().Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// WriteEndpoing
|
||||
// event: endpoint
|
||||
// data: /message?sessionId=285d67ee-1c17-40d9-ab03-173d5ff48419
|
||||
func (w *SSEWriter) WriteEndpoing() {
|
||||
w.c.Writer.WriteString(fmt.Sprintf("event: endpoint\n"))
|
||||
w.c.Writer.WriteString(fmt.Sprintf("data: /message?sessionId=%s\n\n", w.id))
|
||||
w.c.Writer.Flush()
|
||||
}
|
||||
|
||||
// WritePing
|
||||
// : ping - 2025-03-16 06:41:51.280928+00:00
|
||||
func (w *SSEWriter) writePing() {
|
||||
w.c.Writer.WriteString(fmt.Sprintf(": ping - %s\n\n", time.Now().Format("2006-01-02 15:04:05.999999-07:00")))
|
||||
}
|
||||
|
||||
// SSE Session
|
||||
// 维持一个 SSE 连接的会话
|
||||
// 会话中包含了 SSE 连接的 ID，事件通道，停止通道
|
||||
// 事件通道用于发送事件，停止通道用于停止会话
|
||||
// 需要轮询发送 ping 事件以保持连接
|
||||
type SSESession struct {
|
||||
SessionID string
|
||||
Events map[string]chan string
|
||||
Stop chan bool
|
||||
|
||||
c *gin.Context
|
||||
}
|
||||
|
||||
func NewSSESession(c *gin.Context) *SSESession {
|
||||
return &SSESession{c: c}
|
||||
}
|
||||
func (s *SSESession) SendEvent(event string, data string) {
|
||||
s.c.SSEvent(event, data)
|
||||
}
|
||||
|
||||
func (s *SSESession) Close() {
|
||||
close(s.Stop)
|
||||
}
|
||||
|
||||
// Event
|
||||
// request:
|
||||
// POST /messages?sesessionId=?
|
||||
// '{"method":"prompts/list","params":{},"jsonrpc":"2.0","id":3}'
|
||||
//
|
||||
// response:
|
||||
// GET /sse
|
||||
// event: message
|
||||
// data: {"jsonrpc":"2.0","id":3,"result":{"prompts":[]}}
|
||||
|
||||
// {
|
||||
// "jsonrpc": "2.0",
|
||||
// "id": 1,
|
||||
// "result": {
|
||||
// "tools": [
|
||||
// {
|
||||
// "name": "get_alerts",
|
||||
// "description": "Get weather alerts for a US state.\n\n Args:\n state: Two-letter US state code (e.g. CA, NY)\n ",
|
||||
// "inputSchema": {
|
||||
// "properties": {
|
||||
// "state": {
|
||||
// "title": "State",
|
||||
// "type": "string"
|
||||
// }
|
||||
// },
|
||||
// "required": [
|
||||
// "state"
|
||||
// ],
|
||||
// "title": "get_alertsArguments",
|
||||
// "type": "object"
|
||||
// }
|
||||
// },
|
||||
// {
|
||||
// "name": "get_forecast",
|
||||
// "description": "Get weather forecast for a location.\n\n Args:\n latitude: Latitude of the location\n longitude: Longitude of the location\n ",
|
||||
// "inputSchema": {
|
||||
// "properties": {
|
||||
// "latitude": {
|
||||
// "title": "Latitude",
|
||||
// "type": "number"
|
||||
// },
|
||||
// "longitude": {
|
||||
// "title": "Longitude",
|
||||
// "type": "number"
|
||||
// }
|
||||
// },
|
||||
// "required": [
|
||||
// "latitude",
|
||||
// "longitude"
|
||||
// ],
|
||||
// "title": "get_forecastArguments",
|
||||
// "type": "object"
|
||||
// }
|
||||
// }
|
||||
// ]
|
||||
// }
|
||||
// }
|
||||
|
||||
// PING
|
||||
1
internal/mcp/stdio.go
Normal file
1
internal/mcp/stdio.go
Normal file
@@ -0,0 +1 @@
|
||||
package mcp
|
||||
142
internal/mcp/tool.go
Normal file
142
internal/mcp/tool.go
Normal file
@@ -0,0 +1,142 @@
|
||||
package mcp
|
||||
|
||||
// Document: https://modelcontextprotocol.io/docs/concepts/tools
|
||||
|
||||
const (
|
||||
// Client => Server
|
||||
MethodToolsList = "tools/list"
|
||||
MethodToolsCall = "tools/call"
|
||||
)
|
||||
|
||||
type M map[string]interface{}
|
||||
|
||||
// Tool
|
||||
//
|
||||
// {
|
||||
// name: string; // Unique identifier for the tool
|
||||
// description?: string; // Human-readable description
|
||||
// inputSchema: { // JSON Schema for the tool's parameters
|
||||
// type: "object",
|
||||
// properties: { ... } // Tool-specific parameters
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// {
|
||||
// name: "analyze_csv",
|
||||
// description: "Analyze a CSV file",
|
||||
// inputSchema: {
|
||||
// type: "object",
|
||||
// properties: {
|
||||
// filepath: { type: "string" },
|
||||
// operations: {
|
||||
// type: "array",
|
||||
// items: {
|
||||
// enum: ["sum", "average", "count"]
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// {
|
||||
// "jsonrpc": "2.0",
|
||||
// "id": 1,
|
||||
// "result": {
|
||||
// "tools": [
|
||||
// {
|
||||
// "name": "get_alerts",
|
||||
// "description": "Get weather alerts for a US state.\n\n Args:\n state: Two-letter US state code (e.g. CA, NY)\n ",
|
||||
// "inputSchema": {
|
||||
// "properties": {
|
||||
// "state": {
|
||||
// "title": "State",
|
||||
// "type": "string"
|
||||
// }
|
||||
// },
|
||||
// "required": [
|
||||
// "state"
|
||||
// ],
|
||||
// "title": "get_alertsArguments",
|
||||
// "type": "object"
|
||||
// }
|
||||
// },
|
||||
// {
|
||||
// "name": "get_forecast",
|
||||
// "description": "Get weather forecast for a location.\n\n Args:\n latitude: Latitude of the location\n longitude: Longitude of the location\n ",
|
||||
// "inputSchema": {
|
||||
// "properties": {
|
||||
// "latitude": {
|
||||
// "title": "Latitude",
|
||||
// "type": "number"
|
||||
// },
|
||||
// "longitude": {
|
||||
// "title": "Longitude",
|
||||
// "type": "number"
|
||||
// }
|
||||
// },
|
||||
// "required": [
|
||||
// "latitude",
|
||||
// "longitude"
|
||||
// ],
|
||||
// "title": "get_forecastArguments",
|
||||
// "type": "object"
|
||||
// }
|
||||
// }
|
||||
// ]
|
||||
// }
|
||||
// }
|
||||
type Tool struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description,omitempty"`
|
||||
InputSchema ToolSchema `json:"inputSchema"`
|
||||
}
|
||||
|
||||
type ToolSchema struct {
|
||||
Type string `json:"type"`
|
||||
Properties M `json:"properties"`
|
||||
}
|
||||
|
||||
// {
|
||||
// "method": "tools/call",
|
||||
// "params": {
|
||||
// "name": "chatlog",
|
||||
// "arguments": {
|
||||
// "start": "2006-11-12",
|
||||
// "end": "2020-11-20",
|
||||
// "limit": "50",
|
||||
// "offset": "6"
|
||||
// },
|
||||
// "_meta": {
|
||||
// "progressToken": 1
|
||||
// }
|
||||
// },
|
||||
// "jsonrpc": "2.0",
|
||||
// "id": 3
|
||||
// }
|
||||
type ToolsCallRequest struct {
|
||||
Name string `json:"name"`
|
||||
Arguments M `json:"arguments"`
|
||||
}
|
||||
|
||||
// {
|
||||
// "jsonrpc": "2.0",
|
||||
// "id": 2,
|
||||
// "result": {
|
||||
// "content": [
|
||||
// {
|
||||
// "type": "text",
|
||||
// "text": "\nEvent: Winter Storm Warning\n"
|
||||
// }
|
||||
// ],
|
||||
// "isError": false
|
||||
// }
|
||||
// }
|
||||
type ToolsCallResponse struct {
|
||||
Content []Content `json:"content"`
|
||||
IsError bool `json:"isError"`
|
||||
}
|
||||
|
||||
type Content struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text"`
|
||||
}
|
||||
68
internal/ui/footer/footer.go
Normal file
68
internal/ui/footer/footer.go
Normal file
@@ -0,0 +1,68 @@
|
||||
package footer
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/sjzar/chatlog/internal/ui/style"
|
||||
"github.com/sjzar/chatlog/pkg/version"
|
||||
|
||||
"github.com/rivo/tview"
|
||||
)
|
||||
|
||||
const (
|
||||
Title = "footer"
|
||||
)
|
||||
|
||||
type Footer struct {
|
||||
*tview.Flex
|
||||
title string
|
||||
copyRight *tview.TextView
|
||||
help *tview.TextView
|
||||
}
|
||||
|
||||
func New() *Footer {
|
||||
footer := &Footer{
|
||||
Flex: tview.NewFlex(),
|
||||
title: Title,
|
||||
copyRight: tview.NewTextView(),
|
||||
help: tview.NewTextView(),
|
||||
}
|
||||
|
||||
footer.copyRight.
|
||||
SetDynamicColors(true).
|
||||
SetWrap(true).
|
||||
SetTextAlign(tview.AlignLeft)
|
||||
footer.copyRight.
|
||||
SetBackgroundColor(tview.Styles.PrimitiveBackgroundColor)
|
||||
footer.copyRight.SetText(fmt.Sprintf("[%s::b]%s[-:-:-]", style.GetColorHex(style.PageHeaderFgColor), fmt.Sprintf(" @ Sarv's Chatlog (%s)", version.Version)))
|
||||
|
||||
footer.help.
|
||||
SetDynamicColors(true).
|
||||
SetWrap(true).
|
||||
SetTextAlign(tview.AlignRight)
|
||||
footer.help.
|
||||
SetBackgroundColor(tview.Styles.PrimitiveBackgroundColor)
|
||||
|
||||
fmt.Fprintf(footer.help,
|
||||
"[%s::b]↑/↓[%s::b]: 导航 [%s::b]←/→[%s::b]: 切换标签 [%s::b]Enter[%s::b]: 选择 [%s::b]ESC[%s::b]: 返回 [%s::b]Ctrl+C[%s::b]: 退出",
|
||||
style.GetColorHex(style.MenuBgColor), style.GetColorHex(style.PageHeaderFgColor),
|
||||
style.GetColorHex(style.MenuBgColor), style.GetColorHex(style.PageHeaderFgColor),
|
||||
style.GetColorHex(style.MenuBgColor), style.GetColorHex(style.PageHeaderFgColor),
|
||||
style.GetColorHex(style.MenuBgColor), style.GetColorHex(style.PageHeaderFgColor),
|
||||
style.GetColorHex(style.MenuBgColor), style.GetColorHex(style.PageHeaderFgColor),
|
||||
)
|
||||
|
||||
footer.
|
||||
AddItem(footer.copyRight, 0, 1, false).
|
||||
AddItem(footer.help, 0, 1, false)
|
||||
|
||||
return footer
|
||||
}
|
||||
|
||||
func (f *Footer) SetCopyRight(text string) {
|
||||
f.copyRight.SetText(text)
|
||||
}
|
||||
|
||||
func (f *Footer) SetHelp(text string) {
|
||||
f.help.SetText(text)
|
||||
}
|
||||
87
internal/ui/help/help.go
Normal file
87
internal/ui/help/help.go
Normal file
@@ -0,0 +1,87 @@
|
||||
package help
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/sjzar/chatlog/internal/ui/style"
|
||||
|
||||
"github.com/rivo/tview"
|
||||
)
|
||||
|
||||
const (
|
||||
Title = "help"
|
||||
ShowTitle = "帮助"
|
||||
Content = `[yellow]Chatlog 使用指南[white]
|
||||
|
||||
[green]基本操作:[white]
|
||||
• 使用 [yellow]←→[white] 键在主菜单和帮助页面之间切换
|
||||
• 使用 [yellow]↑↓[white] 键在菜单项之间移动
|
||||
• 按 [yellow]Enter[white] 选择菜单项
|
||||
• 按 [yellow]Esc[white] 返回上一级菜单
|
||||
• 按 [yellow]Ctrl+C[white] 退出程序
|
||||
|
||||
[green]使用步骤:[white]
|
||||
|
||||
[yellow]1. 获取数据密钥[white]
|
||||
选择"获取数据密钥"菜单项，程序会自动从运行中的微信进程获取密钥。
|
||||
如果有多个微信进程，会自动选择当前账号的进程。
|
||||
确保微信正在运行，否则无法获取密钥。
|
||||
|
||||
[yellow]2. 解密数据[white]
|
||||
选择"解密数据"菜单项，程序会使用获取的密钥解密微信数据库文件。
|
||||
解密后的文件会保存到工作目录中（可在设置中修改）。
|
||||
|
||||
[yellow]3. 启动 HTTP 服务[white]
|
||||
选择"启动 HTTP 服务"菜单项，启动 HTTP 和 MCP 服务。
|
||||
启动后可以通过浏览器访问 http://localhost:5030 查看聊天记录。
|
||||
|
||||
[yellow]4. 设置选项[white]
|
||||
选择"设置"菜单项，可以配置:
|
||||
• HTTP 服务端口 - 更改 HTTP 服务的监听端口
|
||||
• 工作目录 - 更改解密数据的存储位置
|
||||
|
||||
[green]HTTP API 使用:[white]
|
||||
• 聊天记录: [yellow]GET http://localhost:5030/api/v1/chatlog?time=2023-01-01&talker=wxid_xxx[white]
|
||||
• 联系人列表: [yellow]GET http://localhost:5030/api/v1/contact[white]
|
||||
• 群聊列表: [yellow]GET http://localhost:5030/api/v1/chatroom[white]
|
||||
• 会话列表: [yellow]GET http://localhost:5030/api/v1/session[white]
|
||||
|
||||
[green]MCP 集成:[white]
|
||||
Chatlog 支持 Model Context Protocol，可与支持 MCP 的 AI 助手集成。
|
||||
通过 MCP，AI 助手可以直接查询您的聊天记录、联系人和群聊信息。
|
||||
|
||||
[green]常见问题:[white]
|
||||
• 如果获取密钥失败，请确保微信程序正在运行
|
||||
• 如果解密失败，请检查密钥是否正确获取
|
||||
• 如果 HTTP 服务启动失败，请检查端口是否被占用
|
||||
• 数据目录和工作目录会自动保存，下次启动时自动加载
|
||||
|
||||
[green]数据安全:[white]
|
||||
• 所有数据处理均在本地完成，不会上传到任何外部服务器
|
||||
• 请妥善保管解密后的数据，避免隐私泄露
|
||||
`
|
||||
)
|
||||
|
||||
type Help struct {
|
||||
*tview.TextView
|
||||
title string
|
||||
}
|
||||
|
||||
func New() *Help {
|
||||
help := &Help{
|
||||
TextView: tview.NewTextView(),
|
||||
title: Title,
|
||||
}
|
||||
|
||||
help.SetDynamicColors(true)
|
||||
help.SetRegions(true)
|
||||
help.SetWrap(true)
|
||||
help.SetTextAlign(tview.AlignLeft)
|
||||
help.SetBorder(true)
|
||||
help.SetBorderColor(style.BorderColor)
|
||||
help.SetTitle(ShowTitle)
|
||||
|
||||
fmt.Fprint(help, Content)
|
||||
|
||||
return help
|
||||
}
|
||||
182
internal/ui/infobar/infobar.go
Normal file
182
internal/ui/infobar/infobar.go
Normal file
@@ -0,0 +1,182 @@
|
||||
package infobar
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/sjzar/chatlog/internal/ui/style"
|
||||
|
||||
"github.com/gdamore/tcell/v2"
|
||||
"github.com/rivo/tview"
|
||||
)
|
||||
|
||||
const (
|
||||
Title = "infobar"
|
||||
)
|
||||
|
||||
// InfoBarViewHeight info bar height.
|
||||
const (
|
||||
InfoBarViewHeight = 6
|
||||
accountRow = 0
|
||||
pidRow = 1
|
||||
statusRow = 2
|
||||
dataUsageRow = 3
|
||||
workUsageRow = 4
|
||||
httpServerRow = 5
|
||||
|
||||
// 列索引
|
||||
labelCol1 = 0 // 第一列标签
|
||||
valueCol1 = 1 // 第一列值
|
||||
labelCol2 = 2 // 第二列标签
|
||||
valueCol2 = 3 // 第二列值
|
||||
totalCols = 4
|
||||
)
|
||||
|
||||
// InfoBar implements the info bar primitive.
|
||||
type InfoBar struct {
|
||||
*tview.Box
|
||||
title string
|
||||
table *tview.Table
|
||||
}
|
||||
|
||||
// NewInfoBar returns info bar view.
|
||||
func New() *InfoBar {
|
||||
table := tview.NewTable()
|
||||
headerColor := style.InfoBarItemFgColor
|
||||
|
||||
// Account 和 Version 行
|
||||
table.SetCell(
|
||||
accountRow,
|
||||
labelCol1,
|
||||
tview.NewTableCell(fmt.Sprintf(" [%s::]%s", headerColor, "Account:")),
|
||||
)
|
||||
table.SetCell(accountRow, valueCol1, tview.NewTableCell(""))
|
||||
|
||||
table.SetCell(
|
||||
accountRow,
|
||||
labelCol2,
|
||||
tview.NewTableCell(fmt.Sprintf(" [%s::]%s", headerColor, "Version:")),
|
||||
)
|
||||
table.SetCell(accountRow, valueCol2, tview.NewTableCell(""))
|
||||
|
||||
// PID 和 ExePath 行
|
||||
table.SetCell(
|
||||
pidRow,
|
||||
labelCol1,
|
||||
tview.NewTableCell(fmt.Sprintf(" [%s::]%s", headerColor, "PID:")),
|
||||
)
|
||||
table.SetCell(pidRow, valueCol1, tview.NewTableCell(""))
|
||||
|
||||
table.SetCell(
|
||||
pidRow,
|
||||
labelCol2,
|
||||
tview.NewTableCell(fmt.Sprintf(" [%s::]%s", headerColor, "ExePath:")),
|
||||
)
|
||||
table.SetCell(pidRow, valueCol2, tview.NewTableCell(""))
|
||||
|
||||
// Status 和 Key 行
|
||||
table.SetCell(
|
||||
statusRow,
|
||||
labelCol1,
|
||||
tview.NewTableCell(fmt.Sprintf(" [%s::]%s", headerColor, "Status:")),
|
||||
)
|
||||
table.SetCell(statusRow, valueCol1, tview.NewTableCell(""))
|
||||
|
||||
table.SetCell(
|
||||
statusRow,
|
||||
labelCol2,
|
||||
tview.NewTableCell(fmt.Sprintf(" [%s::]%s", headerColor, "Data Key:")),
|
||||
)
|
||||
table.SetCell(statusRow, valueCol2, tview.NewTableCell(""))
|
||||
|
||||
// Data Usage 和 Data Dir 行
|
||||
table.SetCell(
|
||||
dataUsageRow,
|
||||
labelCol1,
|
||||
tview.NewTableCell(fmt.Sprintf(" [%s::]%s", headerColor, "Data Usage:")),
|
||||
)
|
||||
table.SetCell(dataUsageRow, valueCol1, tview.NewTableCell(""))
|
||||
|
||||
table.SetCell(
|
||||
dataUsageRow,
|
||||
labelCol2,
|
||||
tview.NewTableCell(fmt.Sprintf(" [%s::]%s", headerColor, "Data Dir:")),
|
||||
)
|
||||
table.SetCell(dataUsageRow, valueCol2, tview.NewTableCell(""))
|
||||
|
||||
// Work Usage 和 Work Dir 行
|
||||
table.SetCell(
|
||||
workUsageRow,
|
||||
labelCol1,
|
||||
tview.NewTableCell(fmt.Sprintf(" [%s::]%s", headerColor, "Work Usage:")),
|
||||
)
|
||||
table.SetCell(workUsageRow, valueCol1, tview.NewTableCell(""))
|
||||
|
||||
table.SetCell(
|
||||
workUsageRow,
|
||||
labelCol2,
|
||||
tview.NewTableCell(fmt.Sprintf(" [%s::]%s", headerColor, "Work Dir:")),
|
||||
)
|
||||
table.SetCell(workUsageRow, valueCol2, tview.NewTableCell(""))
|
||||
|
||||
// HTTP Server 行
|
||||
table.SetCell(
|
||||
httpServerRow,
|
||||
labelCol1,
|
||||
tview.NewTableCell(fmt.Sprintf(" [%s::]%s", headerColor, "HTTP Server:")),
|
||||
)
|
||||
table.SetCell(httpServerRow, valueCol1, tview.NewTableCell(""))
|
||||
|
||||
// infobar
|
||||
infoBar := &InfoBar{
|
||||
Box: tview.NewBox(),
|
||||
title: Title,
|
||||
table: table,
|
||||
}
|
||||
|
||||
return infoBar
|
||||
}
|
||||
|
||||
func (info *InfoBar) UpdateAccount(account string) {
|
||||
info.table.GetCell(accountRow, valueCol1).SetText(account)
|
||||
}
|
||||
|
||||
func (info *InfoBar) UpdateBasicInfo(pid int, version string, exePath string) {
|
||||
info.table.GetCell(pidRow, valueCol1).SetText(fmt.Sprintf("%d", pid))
|
||||
info.table.GetCell(pidRow, valueCol2).SetText(exePath)
|
||||
info.table.GetCell(accountRow, valueCol2).SetText(version)
|
||||
}
|
||||
|
||||
func (info *InfoBar) UpdateStatus(status string) {
|
||||
info.table.GetCell(statusRow, valueCol1).SetText(status)
|
||||
}
|
||||
|
||||
func (info *InfoBar) UpdateDataKey(key string) {
|
||||
info.table.GetCell(statusRow, valueCol2).SetText(key)
|
||||
}
|
||||
|
||||
func (info *InfoBar) UpdateDataUsageDir(dataUsage string, dataDir string) {
|
||||
info.table.GetCell(dataUsageRow, valueCol1).SetText(dataUsage)
|
||||
info.table.GetCell(dataUsageRow, valueCol2).SetText(dataDir)
|
||||
}
|
||||
|
||||
func (info *InfoBar) UpdateWorkUsageDir(workUsage string, workDir string) {
|
||||
info.table.GetCell(workUsageRow, valueCol1).SetText(workUsage)
|
||||
info.table.GetCell(workUsageRow, valueCol2).SetText(workDir)
|
||||
}
|
||||
|
||||
// UpdateHTTPServer updates HTTP Server value.
|
||||
func (info *InfoBar) UpdateHTTPServer(server string) {
|
||||
info.table.GetCell(httpServerRow, valueCol1).SetText(server)
|
||||
}
|
||||
|
||||
// Draw draws this primitive onto the screen.
|
||||
func (info *InfoBar) Draw(screen tcell.Screen) {
|
||||
info.Box.DrawForSubclass(screen, info)
|
||||
info.Box.SetBorder(false)
|
||||
|
||||
x, y, width, height := info.GetInnerRect()
|
||||
|
||||
info.table.SetRect(x, y, width, height)
|
||||
info.table.SetBorder(false)
|
||||
info.table.Draw(screen)
|
||||
}
|
||||
162
internal/ui/menu/menu.go
Normal file
162
internal/ui/menu/menu.go
Normal file
@@ -0,0 +1,162 @@
|
||||
package menu
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sort"
|
||||
|
||||
"github.com/sjzar/chatlog/internal/ui/style"
|
||||
|
||||
"github.com/gdamore/tcell/v2"
|
||||
"github.com/rivo/tview"
|
||||
)
|
||||
|
||||
type Item struct {
|
||||
Index int
|
||||
Key string
|
||||
Name string
|
||||
Description string
|
||||
Hidden bool
|
||||
Selected func(i *Item)
|
||||
}
|
||||
|
||||
type Menu struct {
|
||||
*tview.Box
|
||||
title string
|
||||
table *tview.Table
|
||||
items []*Item
|
||||
}
|
||||
|
||||
func New(title string) *Menu {
|
||||
menu := &Menu{
|
||||
Box: tview.NewBox(),
|
||||
title: title,
|
||||
items: make([]*Item, 0),
|
||||
table: tview.NewTable(),
|
||||
}
|
||||
|
||||
menu.table.SetBorders(false)
|
||||
menu.table.SetSelectable(true, false)
|
||||
menu.table.SetTitle(fmt.Sprintf("[::b]%s", menu.title))
|
||||
menu.table.SetBorderColor(style.BorderColor)
|
||||
menu.table.SetBackgroundColor(style.BgColor)
|
||||
menu.table.SetTitleColor(style.FgColor)
|
||||
menu.table.SetFixed(1, 0)
|
||||
menu.table.Select(1, 0).SetSelectedFunc(func(row, column int) {
|
||||
if row == 0 {
|
||||
return // 忽略表头
|
||||
}
|
||||
|
||||
item, ok := menu.table.GetCell(row, 0).GetReference().(*Item)
|
||||
if ok {
|
||||
if item.Selected != nil {
|
||||
item.Selected(item)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
menu.setTableHeader()
|
||||
|
||||
return menu
|
||||
}
|
||||
|
||||
func (m *Menu) setTableHeader() {
|
||||
m.table.SetCell(0, 0, tview.NewTableCell(fmt.Sprintf("[black::b]%s", "命令")).
|
||||
SetExpansion(1).
|
||||
SetBackgroundColor(style.PageHeaderBgColor).
|
||||
SetTextColor(style.PageHeaderFgColor).
|
||||
SetAlign(tview.AlignLeft).
|
||||
SetSelectable(false))
|
||||
|
||||
m.table.SetCell(0, 1, tview.NewTableCell(fmt.Sprintf("[black::b]%s", "说明")).
|
||||
SetExpansion(2).
|
||||
SetBackgroundColor(style.PageHeaderBgColor).
|
||||
SetTextColor(style.PageHeaderFgColor).
|
||||
SetAlign(tview.AlignLeft).
|
||||
SetSelectable(false))
|
||||
}
|
||||
|
||||
func (m *Menu) AddItem(item *Item) {
|
||||
m.items = append(m.items, item)
|
||||
sort.Sort(SortItems(m.items))
|
||||
m.refresh()
|
||||
}
|
||||
|
||||
func (m *Menu) SetItems(items []*Item) {
|
||||
m.items = items
|
||||
m.refresh()
|
||||
}
|
||||
|
||||
func (m *Menu) GetItems() []*Item {
|
||||
return m.items
|
||||
}
|
||||
|
||||
func (m *Menu) refresh() {
|
||||
m.table.Clear()
|
||||
m.setTableHeader()
|
||||
|
||||
row := 1
|
||||
for _, item := range m.items {
|
||||
if item.Hidden {
|
||||
continue
|
||||
}
|
||||
m.table.SetCell(row, 0, tview.NewTableCell(item.Name).
|
||||
SetTextColor(style.FgColor).
|
||||
SetBackgroundColor(style.BgColor).
|
||||
SetReference(item).
|
||||
SetAlign(tview.AlignLeft))
|
||||
m.table.SetCell(row, 1, tview.NewTableCell(item.Description).
|
||||
SetTextColor(style.FgColor).
|
||||
SetBackgroundColor(style.BgColor).
|
||||
SetReference(item).
|
||||
SetAlign(tview.AlignLeft))
|
||||
row++
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
func (m *Menu) Draw(screen tcell.Screen) {
|
||||
m.refresh()
|
||||
|
||||
m.Box.DrawForSubclass(screen, m)
|
||||
m.Box.SetBorder(false)
|
||||
|
||||
menuViewX, menuViewY, menuViewW, menuViewH := m.GetInnerRect()
|
||||
|
||||
m.table.SetRect(menuViewX, menuViewY, menuViewW, menuViewH)
|
||||
m.table.SetBorder(true).SetBorderColor(style.BorderColor)
|
||||
|
||||
m.table.Draw(screen)
|
||||
}
|
||||
|
||||
func (m *Menu) Focus(delegate func(p tview.Primitive)) {
|
||||
delegate(m.table)
|
||||
}
|
||||
|
||||
// HasFocus returns whether or not this primitive has focus
|
||||
func (m *Menu) HasFocus() bool {
|
||||
// Check if the active menu has focus
|
||||
return m.table.HasFocus()
|
||||
}
|
||||
|
||||
func (m *Menu) InputHandler() func(event *tcell.EventKey, setFocus func(p tview.Primitive)) {
|
||||
return m.WrapInputHandler(func(event *tcell.EventKey, setFocus func(p tview.Primitive)) {
|
||||
// 将事件传递给表格
|
||||
if handler := m.table.InputHandler(); handler != nil {
|
||||
handler(event, setFocus)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
type SortItems []*Item
|
||||
|
||||
func (l SortItems) Len() int {
|
||||
return len(l)
|
||||
}
|
||||
|
||||
func (l SortItems) Less(i, j int) bool {
|
||||
return l[i].Index < l[j].Index
|
||||
}
|
||||
|
||||
func (l SortItems) Swap(i, j int) {
|
||||
l[i], l[j] = l[j], l[i]
|
||||
}
|
||||
232
internal/ui/menu/submenu.go
Normal file
232
internal/ui/menu/submenu.go
Normal file
@@ -0,0 +1,232 @@
|
||||
package menu
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sort"
|
||||
|
||||
"github.com/sjzar/chatlog/internal/ui/style"
|
||||
|
||||
"github.com/gdamore/tcell/v2"
|
||||
"github.com/rivo/tview"
|
||||
)
|
||||
|
||||
const (
|
||||
// DialogPadding dialog inner paddign.
|
||||
DialogPadding = 3
|
||||
|
||||
// DialogFormHeight dialog "Enter"/"Cancel" form height.
|
||||
DialogHelpHeight = 1
|
||||
|
||||
// DialogMinWidth dialog min width.
|
||||
DialogMinWidth = 40
|
||||
|
||||
// TableHeightOffset table height offset for border.
|
||||
TableHeightOffset = 3
|
||||
|
||||
cmdWidthOffset = 6
|
||||
)
|
||||
|
||||
type SubMenu struct {
|
||||
*tview.Box
|
||||
title string
|
||||
layout *tview.Flex
|
||||
table *tview.Table
|
||||
width int
|
||||
height int
|
||||
items []*Item
|
||||
cancelHandler func()
|
||||
}
|
||||
|
||||
func NewSubMenu(title string) *SubMenu {
|
||||
subMenu := &SubMenu{
|
||||
Box: tview.NewBox(),
|
||||
title: title,
|
||||
items: make([]*Item, 0),
|
||||
layout: tview.NewFlex(),
|
||||
table: tview.NewTable(),
|
||||
}
|
||||
|
||||
subMenu.table.SetBorders(false)
|
||||
subMenu.table.SetSelectable(true, false)
|
||||
subMenu.table.SetBorderColor(style.DialogBorderColor)
|
||||
subMenu.table.SetBackgroundColor(style.DialogBgColor)
|
||||
subMenu.table.SetTitleColor(style.DialogFgColor)
|
||||
subMenu.table.SetFixed(1, 1)
|
||||
|
||||
subMenu.table.Select(1, 0).SetSelectedFunc(func(row, column int) {
|
||||
if row == 0 {
|
||||
return // 忽略表头
|
||||
}
|
||||
|
||||
item := subMenu.items[row-1]
|
||||
if item.Selected != nil {
|
||||
item.Selected(item)
|
||||
}
|
||||
})
|
||||
|
||||
subMenu.setTableHeader()
|
||||
|
||||
// 帮助信息
|
||||
helpText := tview.NewTextView()
|
||||
helpText.SetDynamicColors(true)
|
||||
helpText.SetTextAlign(tview.AlignCenter)
|
||||
helpText.SetTextColor(style.DialogFgColor)
|
||||
helpText.SetBackgroundColor(style.DialogBgColor)
|
||||
fmt.Fprintf(helpText,
|
||||
"[%s::b]↑/↓[%s::b]: 导航 [%s::b]Enter[%s::b]: 选择 [%s::b]ESC[%s::b]: 返回",
|
||||
style.GetColorHex(style.MenuBgColor), style.GetColorHex(style.PageHeaderFgColor),
|
||||
style.GetColorHex(style.MenuBgColor), style.GetColorHex(style.PageHeaderFgColor),
|
||||
style.GetColorHex(style.MenuBgColor), style.GetColorHex(style.PageHeaderFgColor),
|
||||
)
|
||||
|
||||
// 布局
|
||||
tableLayout := tview.NewFlex().SetDirection(tview.FlexColumn)
|
||||
tableLayout.AddItem(EmptyBoxSpace(style.DialogBgColor), 1, 0, true)
|
||||
tableLayout.AddItem(subMenu.table, 0, 1, true)
|
||||
tableLayout.AddItem(EmptyBoxSpace(style.DialogBgColor), 1, 0, true)
|
||||
|
||||
subMenu.layout.SetDirection(tview.FlexRow)
|
||||
subMenu.layout.SetTitle(fmt.Sprintf("[::b]%s", subMenu.title))
|
||||
subMenu.layout.SetTitleColor(style.DialogFgColor)
|
||||
subMenu.layout.SetTitleAlign(tview.AlignCenter)
|
||||
subMenu.layout.AddItem(tableLayout, 0, 1, true)
|
||||
subMenu.layout.AddItem(helpText, DialogHelpHeight, 0, true)
|
||||
subMenu.layout.SetBorder(true)
|
||||
subMenu.layout.SetBorderColor(style.DialogBorderColor)
|
||||
subMenu.layout.SetBackgroundColor(style.DialogBgColor)
|
||||
|
||||
return subMenu
|
||||
}
|
||||
|
||||
func (m *SubMenu) setTableHeader() {
|
||||
m.table.SetCell(0, 0, tview.NewTableCell(fmt.Sprintf("[%s::b]%s", style.GetColorHex(style.TableHeaderFgColor), "命令")).
|
||||
SetExpansion(1).
|
||||
SetBackgroundColor(style.TableHeaderBgColor).
|
||||
SetTextColor(style.TableHeaderFgColor).
|
||||
SetAlign(tview.AlignLeft).
|
||||
SetSelectable(false))
|
||||
|
||||
m.table.SetCell(0, 1, tview.NewTableCell(fmt.Sprintf("[%s::b]%s", style.GetColorHex(style.TableHeaderFgColor), "说明")).
|
||||
SetExpansion(1).
|
||||
SetBackgroundColor(style.TableHeaderBgColor).
|
||||
SetTextColor(style.TableHeaderFgColor).
|
||||
SetAlign(tview.AlignLeft).
|
||||
SetSelectable(false))
|
||||
}
|
||||
|
||||
func (m *SubMenu) AddItem(item *Item) {
|
||||
m.items = append(m.items, item)
|
||||
sort.Sort(SortItems(m.items))
|
||||
m.refresh()
|
||||
}
|
||||
|
||||
func (m *SubMenu) SetItems(items []*Item) {
|
||||
m.items = items
|
||||
m.refresh()
|
||||
}
|
||||
|
||||
func (m *SubMenu) SetCancelFunc(handler func()) *SubMenu {
|
||||
m.cancelHandler = handler
|
||||
return m
|
||||
}
|
||||
|
||||
func (m *SubMenu) refresh() {
|
||||
m.table.Clear()
|
||||
m.setTableHeader()
|
||||
|
||||
col1Width := 0
|
||||
col2Width := 0
|
||||
|
||||
row := 1
|
||||
for _, item := range m.items {
|
||||
if item.Hidden {
|
||||
continue
|
||||
}
|
||||
m.table.SetCell(row, 0, tview.NewTableCell(item.Name).
|
||||
SetTextColor(style.DialogFgColor).
|
||||
SetBackgroundColor(style.DialogBgColor).
|
||||
SetReference(item).
|
||||
SetAlign(tview.AlignLeft))
|
||||
m.table.SetCell(row, 1, tview.NewTableCell(item.Description).
|
||||
SetTextColor(style.DialogFgColor).
|
||||
SetBackgroundColor(style.DialogBgColor).
|
||||
SetReference(item).
|
||||
SetAlign(tview.AlignLeft))
|
||||
if len(item.Name) > col1Width {
|
||||
col1Width = len(item.Name)
|
||||
}
|
||||
if len(item.Description) > col2Width {
|
||||
col2Width = len(item.Description)
|
||||
}
|
||||
row++
|
||||
}
|
||||
|
||||
m.width = col1Width + col2Width + 2 + cmdWidthOffset
|
||||
m.height = len(m.items) + TableHeightOffset + DialogHelpHeight + 1
|
||||
|
||||
}
|
||||
|
||||
func (m *SubMenu) Draw(screen tcell.Screen) {
|
||||
m.refresh()
|
||||
|
||||
m.Box.DrawForSubclass(screen, m)
|
||||
m.layout.Draw(screen)
|
||||
}
|
||||
|
||||
func (m *SubMenu) SetRect(x, y, width, height int) {
|
||||
ws := (width - m.width) / 2
|
||||
hs := ((height - m.height) / 2)
|
||||
dy := y + hs
|
||||
bWidth := m.width
|
||||
|
||||
if m.width > width {
|
||||
ws = 0
|
||||
bWidth = width - 1
|
||||
}
|
||||
|
||||
bHeight := m.height
|
||||
|
||||
if m.height >= height {
|
||||
dy = y + 1
|
||||
bHeight = height - 1
|
||||
}
|
||||
|
||||
m.Box.SetRect(x+ws, dy, bWidth, bHeight)
|
||||
|
||||
x, y, width, height = m.Box.GetInnerRect()
|
||||
|
||||
m.layout.SetRect(x, y, width, height)
|
||||
}
|
||||
|
||||
func (m *SubMenu) Focus(delegate func(p tview.Primitive)) {
|
||||
delegate(m.table)
|
||||
}
|
||||
|
||||
// HasFocus returns whether or not this primitive has focus
|
||||
func (m *SubMenu) HasFocus() bool {
|
||||
// Check if the active menu has focus
|
||||
return m.table.HasFocus()
|
||||
}
|
||||
|
||||
func (m *SubMenu) InputHandler() func(event *tcell.EventKey, setFocus func(p tview.Primitive)) {
|
||||
return m.WrapInputHandler(func(event *tcell.EventKey, setFocus func(p tview.Primitive)) {
|
||||
|
||||
if event.Key() == tcell.KeyEscape && m.cancelHandler != nil {
|
||||
m.cancelHandler()
|
||||
return
|
||||
}
|
||||
|
||||
// 将事件传递给表格
|
||||
if handler := m.table.InputHandler(); handler != nil {
|
||||
handler(event, setFocus)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func EmptyBoxSpace(bgColor tcell.Color) *tview.Box {
|
||||
box := tview.NewBox()
|
||||
box.SetBackgroundColor(bgColor)
|
||||
box.SetBorder(false)
|
||||
|
||||
return box
|
||||
}
|
||||
78
internal/ui/style/style.go
Normal file
78
internal/ui/style/style.go
Normal file
@@ -0,0 +1,78 @@
|
||||
//go:build !windows
|
||||
// +build !windows
|
||||
|
||||
package style
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/gdamore/tcell/v2"
|
||||
"github.com/rivo/tview"
|
||||
)
|
||||
|
||||
const (
|
||||
// HeavyGreenCheckMark unicode.
|
||||
HeavyGreenCheckMark = "\u2705"
|
||||
// HeavyRedCrossMark unicode.
|
||||
HeavyRedCrossMark = "\u274C"
|
||||
// ProgressBar cell.
|
||||
ProgressBarCell = "▉"
|
||||
)
|
||||
|
||||
var (
|
||||
// infobar.
|
||||
InfoBarItemFgColor = tcell.ColorSilver
|
||||
// main views.
|
||||
FgColor = tcell.ColorFloralWhite
|
||||
BgColor = tview.Styles.PrimitiveBackgroundColor
|
||||
BorderColor = tcell.NewRGBColor(135, 175, 146) //nolint:mnd
|
||||
HelpHeaderFgColor = tcell.NewRGBColor(135, 175, 146) //nolint:mnd
|
||||
MenuBgColor = tcell.ColorMediumSeaGreen
|
||||
PageHeaderBgColor = tcell.ColorMediumSeaGreen
|
||||
PageHeaderFgColor = tcell.ColorFloralWhite
|
||||
RunningStatusFgColor = tcell.NewRGBColor(95, 215, 0) //nolint:mnd
|
||||
PausedStatusFgColor = tcell.NewRGBColor(255, 175, 0) //nolint:mnd
|
||||
// dialogs.
|
||||
DialogBgColor = tcell.NewRGBColor(38, 38, 38) //nolint:mnd
|
||||
DialogBorderColor = tcell.ColorMediumSeaGreen
|
||||
DialogFgColor = tcell.ColorFloralWhite
|
||||
DialogSubBoxBorderColor = tcell.ColorDimGray
|
||||
ErrorDialogBgColor = tcell.NewRGBColor(215, 0, 0) //nolint:mnd
|
||||
ErrorDialogButtonBgColor = tcell.ColorDarkRed
|
||||
// terminal.
|
||||
TerminalFgColor = tcell.ColorFloralWhite
|
||||
TerminalBgColor = tcell.NewRGBColor(5, 5, 5) //nolint:mnd
|
||||
TerminalBorderColor = tcell.ColorDimGray
|
||||
// table header.
|
||||
TableHeaderBgColor = tcell.ColorMediumSeaGreen
|
||||
TableHeaderFgColor = tcell.ColorFloralWhite
|
||||
// progress bar.
|
||||
PrgBgColor = tcell.ColorDimGray
|
||||
PrgBarColor = tcell.ColorDarkOrange
|
||||
PrgBarEmptyColor = tcell.ColorWhite
|
||||
PrgBarOKColor = tcell.ColorGreen
|
||||
PrgBarWarnColor = tcell.ColorOrange
|
||||
PrgBarCritColor = tcell.ColorRed
|
||||
// dropdown.
|
||||
DropDownUnselected = tcell.StyleDefault.Background(tcell.ColorWhiteSmoke).Foreground(tcell.ColorBlack)
|
||||
DropDownSelected = tcell.StyleDefault.Background(tcell.ColorLightSlateGray).Foreground(tcell.ColorWhite)
|
||||
// other primitives.
|
||||
InputFieldBgColor = tcell.ColorGray
|
||||
ButtonBgColor = tcell.ColorMediumSeaGreen
|
||||
)
|
||||
|
||||
// GetColorName returns convert tcell color to its name.
|
||||
func GetColorName(color tcell.Color) string {
|
||||
for name, c := range tcell.ColorNames {
|
||||
if c == color {
|
||||
return name
|
||||
}
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
// GetColorHex returns convert tcell color to its hex useful for textview primitives.
|
||||
func GetColorHex(color tcell.Color) string {
|
||||
return fmt.Sprintf("#%x", color.Hex())
|
||||
}
|
||||
81
internal/ui/style/style_windows.go
Normal file
81
internal/ui/style/style_windows.go
Normal file
@@ -0,0 +1,81 @@
|
||||
//go:build windows
|
||||
|
||||
package style
|
||||
|
||||
import (
|
||||
"github.com/gdamore/tcell/v2"
|
||||
"github.com/rivo/tview"
|
||||
)
|
||||
|
||||
const (
|
||||
// HeavyGreenCheckMark unicode.
|
||||
HeavyGreenCheckMark = "[green::]\u25CF[-::]"
|
||||
// HeavyRedCrossMark unicode.
|
||||
HeavyRedCrossMark = "[red::]\u25CF[-::]"
|
||||
// ProgressBar cell.
|
||||
ProgressBarCell = "\u2593"
|
||||
)
|
||||
|
||||
var (
|
||||
// infobar.
|
||||
InfoBarItemFgColor = tcell.ColorGray
|
||||
// main views.
|
||||
FgColor = tview.Styles.PrimaryTextColor
|
||||
BgColor = tview.Styles.PrimitiveBackgroundColor
|
||||
BorderColor = tcell.ColorSpringGreen
|
||||
MenuBgColor = tcell.ColorSpringGreen
|
||||
HelpHeaderFgColor = tcell.ColorSpringGreen
|
||||
PageHeaderBgColor = tcell.ColorSpringGreen
|
||||
PageHeaderFgColor = tview.Styles.PrimaryTextColor
|
||||
RunningStatusFgColor = tcell.ColorLime
|
||||
PausedStatusFgColor = tcell.ColorYellow
|
||||
|
||||
// dialogs.
|
||||
DialogBgColor = tview.Styles.PrimitiveBackgroundColor
|
||||
DialogFgColor = tview.Styles.PrimaryTextColor
|
||||
DialogBorderColor = tcell.ColorSpringGreen
|
||||
DialogSubBoxBorderColor = tcell.ColorGray
|
||||
ErrorDialogBgColor = tcell.ColorRed
|
||||
ErrorDialogButtonBgColor = tcell.ColorSpringGreen
|
||||
// terminal.
|
||||
TerminalBgColor = tview.Styles.PrimitiveBackgroundColor
|
||||
TerminalFgColor = tview.Styles.PrimaryTextColor
|
||||
TerminalBorderColor = tview.Styles.PrimitiveBackgroundColor
|
||||
// table header.
|
||||
TableHeaderBgColor = tcell.ColorSpringGreen
|
||||
TableHeaderFgColor = tview.Styles.PrimaryTextColor
|
||||
// progress bar.
|
||||
PrgBgColor = tview.Styles.PrimaryTextColor
|
||||
PrgBarColor = tcell.ColorFuchsia
|
||||
PrgBarEmptyColor = tcell.ColorWhite
|
||||
PrgBarOKColor = tcell.ColorLime
|
||||
PrgBarWarnColor = tcell.ColorYellow
|
||||
PrgBarCritColor = tcell.ColorRed
|
||||
// dropdown.
|
||||
DropDownUnselected = tcell.StyleDefault.Background(tcell.ColorGray).Foreground(tcell.ColorWhite)
|
||||
DropDownSelected = tcell.StyleDefault.Background(tcell.ColorPurple).Foreground(tview.Styles.PrimaryTextColor)
|
||||
// other primitives.
|
||||
InputFieldBgColor = tcell.ColorGray
|
||||
ButtonBgColor = tcell.ColorSpringGreen
|
||||
)
|
||||
|
||||
// GetColorName returns convert tcell color to its name.
|
||||
func GetColorName(color tcell.Color) string {
|
||||
for name, c := range tcell.ColorNames {
|
||||
if c == color {
|
||||
return name
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// GetColorHex shall returns convert tcell color to its hex useful for textview primitives,
|
||||
// however, for windows nodes it will return color name.
|
||||
func GetColorHex(color tcell.Color) string {
|
||||
for name, c := range tcell.ColorNames {
|
||||
if c == color {
|
||||
return name
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
415
internal/wechat/decrypt.go
Normal file
415
internal/wechat/decrypt.go
Normal file
@@ -0,0 +1,415 @@
|
||||
package wechat
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/aes"
|
||||
"crypto/cipher"
|
||||
"crypto/hmac"
|
||||
"crypto/sha1"
|
||||
"crypto/sha512"
|
||||
"encoding/binary"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"hash"
|
||||
"io"
|
||||
"os"
|
||||
|
||||
"golang.org/x/crypto/pbkdf2"
|
||||
)
|
||||
|
||||
// Constants for WeChat database decryption
|
||||
const (
|
||||
// Common constants
|
||||
PageSize = 4096
|
||||
KeySize = 32
|
||||
SaltSize = 16
|
||||
AESBlockSize = 16
|
||||
SQLiteHeader = "SQLite format 3\x00"
|
||||
|
||||
// Version specific constants
|
||||
V3IterCount = 64000
|
||||
V4IterCount = 256000
|
||||
|
||||
IVSize = 16 // Same for both versions
|
||||
HmacSHA1Size = 20 // Used in V3
|
||||
HmacSHA512Size = 64 // Used in V4
|
||||
)
|
||||
|
||||
// Error definitions
|
||||
var (
|
||||
ErrHashVerificationFailed = errors.New("hash verification failed")
|
||||
ErrInvalidVersion = errors.New("invalid version, must be 3 or 4")
|
||||
ErrInvalidKey = errors.New("invalid key format")
|
||||
ErrIncorrectKey = errors.New("incorrect decryption key")
|
||||
ErrReadFile = errors.New("failed to read database file")
|
||||
ErrOpenFile = errors.New("failed to open database file")
|
||||
ErrIncompleteRead = errors.New("incomplete header read")
|
||||
ErrCreateCipher = errors.New("failed to create cipher")
|
||||
ErrDecodeKey = errors.New("failed to decode hex key")
|
||||
ErrWriteOutput = errors.New("failed to write output")
|
||||
ErrSeekFile = errors.New("failed to seek in file")
|
||||
ErrOperationCanceled = errors.New("operation was canceled")
|
||||
ErrAlreadyDecrypted = errors.New("file is already decrypted")
|
||||
)
|
||||
|
||||
// Decryptor handles the decryption of WeChat database files
|
||||
type Decryptor struct {
|
||||
// Database file path
|
||||
dbPath string
|
||||
|
||||
// Database properties
|
||||
version int
|
||||
salt []byte
|
||||
page1 []byte
|
||||
reserve int
|
||||
|
||||
// Calculated fields
|
||||
hashFunc func() hash.Hash
|
||||
hmacSize int
|
||||
currentPage int64
|
||||
totalPages int64
|
||||
}
|
||||
|
||||
// NewDecryptor creates a new Decryptor for the specified database file and version
|
||||
func NewDecryptor(dbPath string, version int) (*Decryptor, error) {
|
||||
// Validate version
|
||||
if version != 3 && version != 4 {
|
||||
return nil, ErrInvalidVersion
|
||||
}
|
||||
|
||||
// Open database file
|
||||
fp, err := os.Open(dbPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrOpenFile, err)
|
||||
}
|
||||
defer fp.Close()
|
||||
|
||||
// Get file size
|
||||
fileInfo, err := fp.Stat()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get file info: %v", err)
|
||||
}
|
||||
|
||||
// Calculate total pages
|
||||
fileSize := fileInfo.Size()
|
||||
totalPages := fileSize / PageSize
|
||||
if fileSize%PageSize > 0 {
|
||||
totalPages++
|
||||
}
|
||||
|
||||
// Read first page
|
||||
buffer := make([]byte, PageSize)
|
||||
n, err := io.ReadFull(fp, buffer)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrReadFile, err)
|
||||
}
|
||||
if n != PageSize {
|
||||
return nil, fmt.Errorf("%w: expected %d bytes, got %d", ErrIncompleteRead, PageSize, n)
|
||||
}
|
||||
|
||||
// Check if file is already decrypted
|
||||
if bytes.Equal(buffer[:len(SQLiteHeader)-1], []byte(SQLiteHeader[:len(SQLiteHeader)-1])) {
|
||||
return nil, ErrAlreadyDecrypted
|
||||
}
|
||||
|
||||
// Initialize hash function and HMAC size based on version
|
||||
var hashFunc func() hash.Hash
|
||||
var hmacSize int
|
||||
|
||||
if version == 4 {
|
||||
hashFunc = sha512.New
|
||||
hmacSize = HmacSHA512Size
|
||||
} else {
|
||||
hashFunc = sha1.New
|
||||
hmacSize = HmacSHA1Size
|
||||
}
|
||||
|
||||
// Calculate reserve size and MAC offset
|
||||
reserve := IVSize + hmacSize
|
||||
if reserve%AESBlockSize != 0 {
|
||||
reserve = ((reserve / AESBlockSize) + 1) * AESBlockSize
|
||||
}
|
||||
|
||||
return &Decryptor{
|
||||
dbPath: dbPath,
|
||||
version: version,
|
||||
salt: buffer[:SaltSize],
|
||||
page1: buffer,
|
||||
reserve: reserve,
|
||||
hashFunc: hashFunc,
|
||||
hmacSize: hmacSize,
|
||||
totalPages: totalPages,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// GetTotalPages returns the total number of pages in the database
|
||||
func (d *Decryptor) GetTotalPages() int64 {
|
||||
return d.totalPages
|
||||
}
|
||||
|
||||
// Validate checks if the provided key is valid for this database
|
||||
func (d *Decryptor) Validate(key []byte) bool {
|
||||
if len(key) != KeySize {
|
||||
return false
|
||||
}
|
||||
_, macKey := d.calcPBKDF2Key(key)
|
||||
return d.validate(macKey)
|
||||
}
|
||||
|
||||
func (d *Decryptor) calcPBKDF2Key(key []byte) ([]byte, []byte) {
|
||||
// Generate encryption key from password
|
||||
var encKey []byte
|
||||
if d.version == 4 {
|
||||
encKey = pbkdf2.Key(key, d.salt, V4IterCount, KeySize, sha512.New)
|
||||
} else {
|
||||
encKey = pbkdf2.Key(key, d.salt, V3IterCount, KeySize, sha1.New)
|
||||
}
|
||||
|
||||
// Generate MAC key
|
||||
macSalt := xorBytes(d.salt, 0x3a)
|
||||
macKey := pbkdf2.Key(encKey, macSalt, 2, KeySize, d.hashFunc)
|
||||
return encKey, macKey
|
||||
}
|
||||
|
||||
func (d *Decryptor) validate(macKey []byte) bool {
|
||||
// Calculate HMAC
|
||||
hashMac := hmac.New(d.hashFunc, macKey)
|
||||
|
||||
dataEnd := PageSize - d.reserve + IVSize
|
||||
hashMac.Write(d.page1[SaltSize:dataEnd])
|
||||
|
||||
// Page number is fixed as 1
|
||||
pageNoBytes := make([]byte, 4)
|
||||
binary.LittleEndian.PutUint32(pageNoBytes, 1)
|
||||
hashMac.Write(pageNoBytes)
|
||||
|
||||
calculatedMAC := hashMac.Sum(nil)
|
||||
storedMAC := d.page1[dataEnd : dataEnd+d.hmacSize]
|
||||
|
||||
return hmac.Equal(calculatedMAC, storedMAC)
|
||||
}
|
||||
|
||||
// Decrypt decrypts the database using the provided key and writes the result to the writer
|
||||
func (d *Decryptor) Decrypt(ctx context.Context, hexKey string, w io.Writer) error {
|
||||
// Decode key
|
||||
key, err := hex.DecodeString(hexKey)
|
||||
if err != nil {
|
||||
return fmt.Errorf("%w: %v", ErrDecodeKey, err)
|
||||
}
|
||||
|
||||
encKey, macKey := d.calcPBKDF2Key(key)
|
||||
|
||||
// Validate key first
|
||||
if !d.validate(macKey) {
|
||||
return ErrIncorrectKey
|
||||
}
|
||||
|
||||
// Open input file
|
||||
dbFile, err := os.Open(d.dbPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("%w: %v", ErrOpenFile, err)
|
||||
}
|
||||
defer dbFile.Close()
|
||||
|
||||
// Write SQLite header to output
|
||||
_, err = w.Write([]byte(SQLiteHeader))
|
||||
if err != nil {
|
||||
return fmt.Errorf("%w: %v", ErrWriteOutput, err)
|
||||
}
|
||||
|
||||
// Process each page
|
||||
pageBuf := make([]byte, PageSize)
|
||||
d.currentPage = 0
|
||||
|
||||
for curPage := int64(0); curPage < d.totalPages; curPage++ {
|
||||
// Check for cancellation before processing each page
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ErrOperationCanceled
|
||||
default:
|
||||
// Continue processing
|
||||
}
|
||||
|
||||
// For the first page, we need to skip the salt
|
||||
if curPage == 0 {
|
||||
// Read the first page
|
||||
_, err = io.ReadFull(dbFile, pageBuf)
|
||||
if err != nil {
|
||||
return fmt.Errorf("%w: %v", ErrReadFile, err)
|
||||
}
|
||||
} else {
|
||||
// Read a full page
|
||||
n, err := io.ReadFull(dbFile, pageBuf)
|
||||
if err != nil {
|
||||
if err == io.EOF || err == io.ErrUnexpectedEOF {
|
||||
// Handle last partial page
|
||||
if n > 0 {
|
||||
// Process partial page
|
||||
// For simplicity, we'll just break here
|
||||
break
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("%w: %v", ErrReadFile, err)
|
||||
}
|
||||
}
|
||||
|
||||
// check if page contains only zeros (v3 & v4 both have this behavior)
|
||||
allZeros := true
|
||||
for _, b := range pageBuf {
|
||||
if b != 0 {
|
||||
allZeros = false
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if allZeros {
|
||||
// Write the zeros page to output
|
||||
_, err = w.Write(pageBuf)
|
||||
if err != nil {
|
||||
return fmt.Errorf("%w: %v", ErrWriteOutput, err)
|
||||
}
|
||||
|
||||
// Update progress
|
||||
d.currentPage = curPage + 1
|
||||
continue
|
||||
|
||||
// // Set current page to total pages to indicate completion
|
||||
// d.currentPage = d.totalPages
|
||||
// return nil
|
||||
}
|
||||
|
||||
// Decrypt the page
|
||||
decryptedPage, err := d.decryptPage(encKey, macKey, pageBuf, curPage)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Write decrypted page to output
|
||||
_, err = w.Write(decryptedPage)
|
||||
if err != nil {
|
||||
return fmt.Errorf("%w: %v", ErrWriteOutput, err)
|
||||
}
|
||||
|
||||
// Update progress
|
||||
d.currentPage = curPage + 1
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// decryptPage decrypts a single page of the database
|
||||
func (d *Decryptor) decryptPage(key, macKey []byte, pageBuf []byte, pageNum int64) ([]byte, error) {
|
||||
offset := 0
|
||||
if pageNum == 0 {
|
||||
offset = SaltSize
|
||||
}
|
||||
|
||||
// Verify HMAC
|
||||
mac := hmac.New(d.hashFunc, macKey)
|
||||
mac.Write(pageBuf[offset : PageSize-d.reserve+IVSize])
|
||||
|
||||
// Convert page number and update HMAC
|
||||
pageNumBytes := make([]byte, 4)
|
||||
binary.LittleEndian.PutUint32(pageNumBytes, uint32(pageNum+1))
|
||||
mac.Write(pageNumBytes)
|
||||
|
||||
hashMac := mac.Sum(nil)
|
||||
|
||||
hashMacStartOffset := PageSize - d.reserve + IVSize
|
||||
hashMacEndOffset := hashMacStartOffset + len(hashMac)
|
||||
|
||||
if !bytes.Equal(hashMac, pageBuf[hashMacStartOffset:hashMacEndOffset]) {
|
||||
return nil, ErrHashVerificationFailed
|
||||
}
|
||||
|
||||
// Decrypt content using AES-256-CBC
|
||||
iv := pageBuf[PageSize-d.reserve : PageSize-d.reserve+IVSize]
|
||||
block, err := aes.NewCipher(key)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrCreateCipher, err)
|
||||
}
|
||||
|
||||
mode := cipher.NewCBCDecrypter(block, iv)
|
||||
|
||||
// Create a copy of encrypted data for decryption
|
||||
encrypted := make([]byte, PageSize-d.reserve-offset)
|
||||
copy(encrypted, pageBuf[offset:PageSize-d.reserve])
|
||||
|
||||
// Decrypt in place
|
||||
mode.CryptBlocks(encrypted, encrypted)
|
||||
|
||||
// Combine decrypted data with reserve part
|
||||
decryptedPage := append(encrypted, pageBuf[PageSize-d.reserve:PageSize]...)
|
||||
|
||||
return decryptedPage, nil
|
||||
}
|
||||
|
||||
// xorBytes performs XOR operation on each byte of the array with the specified byte
|
||||
func xorBytes(a []byte, b byte) []byte {
|
||||
result := make([]byte, len(a))
|
||||
for i := range a {
|
||||
result[i] = a[i] ^ b
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// Utility functions for backward compatibility
|
||||
|
||||
// DecryptDBFile decrypts a WeChat database file and returns the decrypted content
|
||||
func DecryptDBFile(dbPath string, hexKey string, version int) ([]byte, error) {
|
||||
// Create a buffer to store the decrypted content
|
||||
var buf bytes.Buffer
|
||||
|
||||
// Create a decryptor
|
||||
d, err := NewDecryptor(dbPath, version)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Decrypt the database
|
||||
err = d.Decrypt(context.Background(), hexKey, &buf)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return buf.Bytes(), nil
|
||||
}
|
||||
|
||||
// DecryptDBFileToFile decrypts a WeChat database file and saves the result to the specified output file
|
||||
func DecryptDBFileToFile(dbPath, outputPath, hexKey string, version int) error {
|
||||
// Create output file
|
||||
outputFile, err := os.Create(outputPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create output file: %v", err)
|
||||
}
|
||||
defer outputFile.Close()
|
||||
|
||||
// Create a decryptor
|
||||
d, err := NewDecryptor(dbPath, version)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Decrypt the database
|
||||
return d.Decrypt(context.Background(), hexKey, outputFile)
|
||||
}
|
||||
|
||||
// ValidateDBKey validates if the provided key is correct for the database
|
||||
func ValidateDBKey(dbPath string, hexKey string, version int) bool {
|
||||
// Create a decryptor
|
||||
d, err := NewDecryptor(dbPath, version)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
|
||||
// Decode key
|
||||
key, err := hex.DecodeString(hexKey)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
|
||||
// Validate the key
|
||||
return d.Validate(key)
|
||||
}
|
||||
50
internal/wechat/info.go
Normal file
50
internal/wechat/info.go
Normal file
@@ -0,0 +1,50 @@
|
||||
package wechat
|
||||
|
||||
import (
|
||||
"github.com/sjzar/chatlog/pkg/dllver"
|
||||
|
||||
"github.com/shirou/gopsutil/v4/process"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
const (
|
||||
StatusInit = ""
|
||||
StatusOffline = "offline"
|
||||
StatusOnline = "online"
|
||||
)
|
||||
|
||||
type Info struct {
|
||||
PID uint32
|
||||
ExePath string
|
||||
Version *dllver.Info
|
||||
Status string
|
||||
DataDir string
|
||||
AccountName string
|
||||
Key string
|
||||
}
|
||||
|
||||
func NewInfo(p *process.Process) (*Info, error) {
|
||||
info := &Info{
|
||||
PID: uint32(p.Pid),
|
||||
Status: StatusOffline,
|
||||
}
|
||||
|
||||
var err error
|
||||
info.ExePath, err = p.Exe()
|
||||
if err != nil {
|
||||
log.Error(err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
info.Version, err = dllver.New(info.ExePath)
|
||||
if err != nil {
|
||||
log.Error(err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if err := info.initialize(p); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return info, nil
|
||||
}
|
||||
16
internal/wechat/info_others.go
Normal file
16
internal/wechat/info_others.go
Normal file
@@ -0,0 +1,16 @@
|
||||
//go:build !windows
|
||||
|
||||
package wechat
|
||||
|
||||
import "github.com/shirou/gopsutil/v4/process"
|
||||
|
||||
// Giao~
|
||||
// 还没来得及写，Mac 版本打算通过 vmmap 检查内存区域，再用 lldb 读取内存来检查 Key，需要关 SIP 或自签名应用，稍晚再填坑
|
||||
|
||||
func (i *Info) initialize(p *process.Process) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (i *Info) GetKey() (string, error) {
|
||||
return "mock-key", nil
|
||||
}
|
||||
507
internal/wechat/info_windows.go
Normal file
507
internal/wechat/info_windows.go
Normal file
@@ -0,0 +1,507 @@
|
||||
package wechat
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
"sync"
|
||||
"unsafe"
|
||||
|
||||
"github.com/sjzar/chatlog/pkg/util"
|
||||
|
||||
"github.com/shirou/gopsutil/v4/process"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"golang.org/x/sys/windows"
|
||||
)
|
||||
|
||||
const (
|
||||
V3ModuleName = "WeChatWin.dll"
|
||||
V3DBFile = "Msg\\Misc.db"
|
||||
V4DBFile = "db_storage\\message\\message_0.db"
|
||||
|
||||
MaxWorkers = 16
|
||||
|
||||
// Windows memory protection constants
|
||||
MEM_PRIVATE = 0x20000
|
||||
)
|
||||
|
||||
// Common error definitions
|
||||
var (
|
||||
ErrWeChatOffline = errors.New("wechat is not logged in")
|
||||
ErrOpenProcess = errors.New("failed to open process")
|
||||
ErrReaddecryptor = errors.New("failed to read database header")
|
||||
ErrCheckProcessBits = errors.New("failed to check process architecture")
|
||||
ErrFindWeChatDLL = errors.New("WeChatWin.dll module not found")
|
||||
ErrNoValidKey = errors.New("no valid key found")
|
||||
ErrInvalidFilePath = errors.New("invalid file path format")
|
||||
)
|
||||
|
||||
// GetKey is the entry point for retrieving the WeChat database key
|
||||
func (i *Info) GetKey() (string, error) {
|
||||
if i.Status == StatusOffline {
|
||||
return "", ErrWeChatOffline
|
||||
}
|
||||
|
||||
// Choose key retrieval method based on WeChat version
|
||||
if i.Version.FileMajorVersion == 4 {
|
||||
return i.getKeyV4()
|
||||
}
|
||||
return i.getKeyV3()
|
||||
}
|
||||
|
||||
// initialize initializes WeChat information
|
||||
func (i *Info) initialize(p *process.Process) error {
|
||||
files, err := p.OpenFiles()
|
||||
if err != nil {
|
||||
log.Error("Failed to get open file list: ", err)
|
||||
return err
|
||||
}
|
||||
|
||||
dbPath := V3DBFile
|
||||
if i.Version.FileMajorVersion == 4 {
|
||||
dbPath = V4DBFile
|
||||
}
|
||||
|
||||
for _, f := range files {
|
||||
if strings.HasSuffix(f.Path, dbPath) {
|
||||
filePath := f.Path[4:] // Remove "\\?\" prefix
|
||||
parts := strings.Split(filePath, string(filepath.Separator))
|
||||
if len(parts) < 4 {
|
||||
log.Debug("Invalid file path format: " + filePath)
|
||||
continue
|
||||
}
|
||||
|
||||
i.Status = StatusOnline
|
||||
if i.Version.FileMajorVersion == 4 {
|
||||
i.DataDir = strings.Join(parts[:len(parts)-3], string(filepath.Separator))
|
||||
i.AccountName = parts[len(parts)-4]
|
||||
} else {
|
||||
i.DataDir = strings.Join(parts[:len(parts)-2], string(filepath.Separator))
|
||||
i.AccountName = parts[len(parts)-3]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// getKeyV3 retrieves the database key for WeChat V3 version
|
||||
func (i *Info) getKeyV3() (string, error) {
|
||||
// Read database header for key validation
|
||||
dbPath := filepath.Join(i.DataDir, V3DBFile)
|
||||
decryptor, err := NewDecryptor(dbPath, i.Version.FileMajorVersion)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w: %v", ErrReaddecryptor, err)
|
||||
}
|
||||
log.Debug("V3 database path: ", dbPath)
|
||||
|
||||
// Open WeChat process
|
||||
handle, err := windows.OpenProcess(windows.PROCESS_QUERY_INFORMATION|windows.PROCESS_VM_READ, false, i.PID)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w: %v", ErrOpenProcess, err)
|
||||
}
|
||||
defer windows.CloseHandle(handle)
|
||||
|
||||
// Check process architecture
|
||||
is64Bit, err := util.Is64Bit(handle)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w: %v", ErrCheckProcessBits, err)
|
||||
}
|
||||
|
||||
// Create context to control all goroutines
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
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
|
||||
}
|
||||
log.Debug("Starting ", workerCount, " workers for V3 key search")
|
||||
|
||||
// Start consumer goroutines
|
||||
var workerWaitGroup sync.WaitGroup
|
||||
workerWaitGroup.Add(workerCount)
|
||||
for index := 0; index < workerCount; index++ {
|
||||
go func() {
|
||||
defer workerWaitGroup.Done()
|
||||
workerV3(ctx, handle, decryptor, is64Bit, 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 := i.findMemoryV3(ctx, handle, memoryChannel)
|
||||
if err != nil {
|
||||
log.Error(err)
|
||||
}
|
||||
}()
|
||||
|
||||
// Wait for producer and consumers to complete
|
||||
go func() {
|
||||
producerWaitGroup.Wait()
|
||||
workerWaitGroup.Wait()
|
||||
close(resultChannel)
|
||||
}()
|
||||
|
||||
// Wait for result
|
||||
result, ok := <-resultChannel
|
||||
if ok && result != "" {
|
||||
i.Key = result
|
||||
return result, nil
|
||||
}
|
||||
|
||||
return "", ErrNoValidKey
|
||||
}
|
||||
|
||||
// getKeyV4 retrieves the database key for WeChat V4 version
|
||||
func (i *Info) getKeyV4() (string, error) {
|
||||
// Read database header for key validation
|
||||
dbPath := filepath.Join(i.DataDir, V4DBFile)
|
||||
decryptor, err := NewDecryptor(dbPath, i.Version.FileMajorVersion)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w: %v", ErrReaddecryptor, err)
|
||||
}
|
||||
log.Debug("V4 database path: ", dbPath)
|
||||
|
||||
// Open process handle
|
||||
handle, err := windows.OpenProcess(windows.PROCESS_VM_READ|windows.PROCESS_QUERY_INFORMATION, false, i.PID)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w: %v", ErrOpenProcess, err)
|
||||
}
|
||||
defer windows.CloseHandle(handle)
|
||||
|
||||
// Create context to control all goroutines
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
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
|
||||
}
|
||||
log.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()
|
||||
workerV4(ctx, handle, decryptor, 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 := i.findMemoryV4(ctx, handle, memoryChannel)
|
||||
if err != nil {
|
||||
log.Error(err)
|
||||
}
|
||||
}()
|
||||
|
||||
// Wait for producer and consumers to complete
|
||||
go func() {
|
||||
producerWaitGroup.Wait()
|
||||
workerWaitGroup.Wait()
|
||||
close(resultChannel)
|
||||
}()
|
||||
|
||||
// Wait for result
|
||||
result, ok := <-resultChannel
|
||||
if ok && result != "" {
|
||||
i.Key = result
|
||||
return result, nil
|
||||
}
|
||||
|
||||
return "", ErrNoValidKey
|
||||
}
|
||||
|
||||
// findMemoryV3 searches for writable memory regions in WeChatWin.dll for V3 version
|
||||
func (i *Info) findMemoryV3(ctx context.Context, handle windows.Handle, memoryChannel chan<- []byte) error {
|
||||
// Find WeChatWin.dll module
|
||||
module, isFound := FindModule(i.PID, V3ModuleName)
|
||||
if !isFound {
|
||||
return ErrFindWeChatDLL
|
||||
}
|
||||
log.Debug("Found WeChatWin.dll module at base address: 0x", fmt.Sprintf("%X", module.ModBaseAddr))
|
||||
|
||||
// Read writable memory regions
|
||||
baseAddr := uintptr(module.ModBaseAddr)
|
||||
endAddr := baseAddr + uintptr(module.ModBaseSize)
|
||||
currentAddr := baseAddr
|
||||
|
||||
for currentAddr < endAddr {
|
||||
var mbi windows.MemoryBasicInformation
|
||||
err := windows.VirtualQueryEx(handle, currentAddr, &mbi, unsafe.Sizeof(mbi))
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
|
||||
// Skip small memory regions
|
||||
if mbi.RegionSize < 100*1024 {
|
||||
currentAddr += uintptr(mbi.RegionSize)
|
||||
continue
|
||||
}
|
||||
|
||||
// Check if memory region is writable
|
||||
isWritable := (mbi.Protect & (windows.PAGE_READWRITE | windows.PAGE_WRITECOPY | windows.PAGE_EXECUTE_READWRITE | windows.PAGE_EXECUTE_WRITECOPY)) > 0
|
||||
if isWritable && uint32(mbi.State) == windows.MEM_COMMIT {
|
||||
// Calculate region size, ensure it doesn't exceed DLL bounds
|
||||
regionSize := uintptr(mbi.RegionSize)
|
||||
if currentAddr+regionSize > endAddr {
|
||||
regionSize = endAddr - currentAddr
|
||||
}
|
||||
|
||||
// Read writable memory region
|
||||
memory := make([]byte, regionSize)
|
||||
if err = windows.ReadProcessMemory(handle, currentAddr, &memory[0], regionSize, nil); err == nil {
|
||||
select {
|
||||
case memoryChannel <- memory:
|
||||
log.Debug("Sent memory region for analysis, size: ", regionSize, " bytes")
|
||||
case <-ctx.Done():
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Move to next memory region
|
||||
currentAddr = uintptr(mbi.BaseAddress) + uintptr(mbi.RegionSize)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// workerV3 processes memory regions to find V3 version key
|
||||
func workerV3(ctx context.Context, handle windows.Handle, decryptor *Decryptor, is64Bit bool, memoryChannel <-chan []byte, resultChannel chan<- string) {
|
||||
|
||||
// Define search pattern
|
||||
keyPattern := []byte{0x20, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
|
||||
ptrSize := 8
|
||||
littleEndianFunc := binary.LittleEndian.Uint64
|
||||
|
||||
// Adjust for 32-bit process
|
||||
if !is64Bit {
|
||||
keyPattern = keyPattern[:4]
|
||||
ptrSize = 4
|
||||
littleEndianFunc = func(b []byte) uint64 { return uint64(binary.LittleEndian.Uint32(b)) }
|
||||
}
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case memory, ok := <-memoryChannel:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
index := len(memory)
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return // Exit if result found
|
||||
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 := validateKey(handle, decryptor, ptrValue); key != "" {
|
||||
select {
|
||||
case resultChannel <- key:
|
||||
log.Debug("Valid key found for V3 database")
|
||||
default:
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
index -= 1 // Continue searching from previous position
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// findMemoryV4 searches for writable memory regions for V4 version
|
||||
func (i *Info) findMemoryV4(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
|
||||
}
|
||||
log.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:
|
||||
log.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 workerV4(ctx context.Context, handle windows.Handle, decryptor *Decryptor, 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 result found
|
||||
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 := validateKey(handle, decryptor, ptrValue); key != "" {
|
||||
select {
|
||||
case resultChannel <- key:
|
||||
log.Debug("Valid key found for V4 database")
|
||||
default:
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
index -= 1 // Continue searching from previous position
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// validateKey validates a single key candidate
|
||||
func validateKey(handle windows.Handle, decryptor *Decryptor, 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 decryptor.Validate(keyData) {
|
||||
return hex.EncodeToString(keyData)
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
// FindModule searches for a specified module in the process
|
||||
// Used to find WeChatWin.dll module for V3 version
|
||||
func FindModule(pid uint32, name string) (module windows.ModuleEntry32, isFound bool) {
|
||||
// Create module snapshot
|
||||
snapshot, err := windows.CreateToolhelp32Snapshot(windows.TH32CS_SNAPMODULE|windows.TH32CS_SNAPMODULE32, pid)
|
||||
if err != nil {
|
||||
log.Debug("Failed to create module snapshot: ", err)
|
||||
return module, false
|
||||
}
|
||||
defer windows.CloseHandle(snapshot)
|
||||
|
||||
// Initialize module entry structure
|
||||
module.Size = uint32(windows.SizeofModuleEntry32)
|
||||
|
||||
// Get the first module
|
||||
if err := windows.Module32First(snapshot, &module); err != nil {
|
||||
log.Debug("Failed to get first module: ", err)
|
||||
return module, false
|
||||
}
|
||||
|
||||
// Iterate through all modules to find WeChatWin.dll
|
||||
for ; err == nil; err = windows.Module32Next(snapshot, &module) {
|
||||
if windows.UTF16ToString(module.Module[:]) == name {
|
||||
return module, true
|
||||
}
|
||||
}
|
||||
return module, false
|
||||
}
|
||||
57
internal/wechat/manager.go
Normal file
57
internal/wechat/manager.go
Normal file
@@ -0,0 +1,57 @@
|
||||
package wechat
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/shirou/gopsutil/v4/process"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
const (
|
||||
V3ProcessName = "WeChat"
|
||||
V4ProcessName = "Weixin"
|
||||
)
|
||||
|
||||
var (
|
||||
Items []*Info
|
||||
ItemMap map[string]*Info
|
||||
)
|
||||
|
||||
func Load() {
|
||||
Items = make([]*Info, 0, 2)
|
||||
ItemMap = make(map[string]*Info)
|
||||
|
||||
processes, err := process.Processes()
|
||||
if err != nil {
|
||||
log.Println("获取进程列表失败:", err)
|
||||
return
|
||||
}
|
||||
|
||||
for _, p := range processes {
|
||||
name, err := p.Name()
|
||||
name = strings.TrimSuffix(name, ".exe")
|
||||
if err != nil || name != V3ProcessName && name != V4ProcessName {
|
||||
continue
|
||||
}
|
||||
|
||||
// v4 存在同名进程，需要继续判断 cmdline
|
||||
if name == V4ProcessName {
|
||||
cmdline, err := p.Cmdline()
|
||||
if err != nil {
|
||||
log.Error(err)
|
||||
continue
|
||||
}
|
||||
if strings.Contains(cmdline, "--") {
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
info, err := NewInfo(p)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
Items = append(Items, info)
|
||||
ItemMap[info.AccountName] = info
|
||||
}
|
||||
}
|
||||
269
internal/wechatdb/contact.go
Normal file
269
internal/wechatdb/contact.go
Normal file
@@ -0,0 +1,269 @@
|
||||
package wechatdb
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
|
||||
"github.com/sjzar/chatlog/pkg/model"
|
||||
"github.com/sjzar/chatlog/pkg/util"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
var (
|
||||
ContactFileV3 = "^MicroMsg.db$"
|
||||
ContactFileV4 = "contact.db$"
|
||||
)
|
||||
|
||||
type Contact struct {
|
||||
version int
|
||||
dbFile string
|
||||
db *sql.DB
|
||||
|
||||
Contact map[string]*model.Contact // 好友和群聊信息，Key UserName
|
||||
ChatRoom map[string]*model.ChatRoom // 群聊信息，Key UserName
|
||||
Sessions []*model.Session // 历史会话，按时间倒序
|
||||
|
||||
// Quick Search
|
||||
ChatRoomUsers map[string]*model.Contact // 群聊成员信息，Key UserName
|
||||
Alias2Contack map[string]*model.Contact // 别名到联系人的映射
|
||||
Remark2Contack map[string]*model.Contact // 备注名到联系人的映射
|
||||
NickName2Contack map[string]*model.Contact // 昵称到联系人的映射
|
||||
}
|
||||
|
||||
func NewContact(path string, version int) (*Contact, error) {
|
||||
c := &Contact{
|
||||
version: version,
|
||||
}
|
||||
|
||||
files, err := util.FindFilesWithPatterns(path, ContactFileV3, true)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查找数据库文件失败: %v", err)
|
||||
}
|
||||
|
||||
if len(files) == 0 {
|
||||
return nil, fmt.Errorf("未找到任何数据库文件: %s", path)
|
||||
}
|
||||
|
||||
c.dbFile = files[0]
|
||||
|
||||
c.db, err = sql.Open("sqlite3", c.dbFile)
|
||||
if err != nil {
|
||||
log.Printf("警告: 连接数据库 %s 失败: %v", c.dbFile, err)
|
||||
return nil, fmt.Errorf("连接数据库失败: %v", err)
|
||||
}
|
||||
|
||||
c.loadContact()
|
||||
c.loadChatRoom()
|
||||
c.loadSession()
|
||||
c.fillChatRoomInfo()
|
||||
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (c *Contact) loadContact() {
|
||||
contactMap := make(map[string]*model.Contact)
|
||||
chatRoomUserMap := make(map[string]*model.Contact)
|
||||
aliasMap := make(map[string]*model.Contact)
|
||||
remarkMap := make(map[string]*model.Contact)
|
||||
nickNameMap := make(map[string]*model.Contact)
|
||||
rows, err := c.db.Query("SELECT UserName, Alias, Remark, NickName, Reserved1 FROM Contact")
|
||||
if err != nil {
|
||||
log.Errorf("查询联系人失败: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
for rows.Next() {
|
||||
var contactv3 model.ContactV3
|
||||
|
||||
if err := rows.Scan(
|
||||
&contactv3.UserName,
|
||||
&contactv3.Alias,
|
||||
&contactv3.Remark,
|
||||
&contactv3.NickName,
|
||||
&contactv3.Reserved1,
|
||||
); err != nil {
|
||||
log.Printf("警告: 扫描联系人行失败: %v", err)
|
||||
continue
|
||||
}
|
||||
contact := contactv3.Wrap()
|
||||
|
||||
if contact.IsFriend {
|
||||
contactMap[contact.UserName] = contact
|
||||
if contact.Alias != "" {
|
||||
aliasMap[contact.Alias] = contact
|
||||
}
|
||||
if contact.Remark != "" {
|
||||
remarkMap[contact.Remark] = contact
|
||||
}
|
||||
if contact.NickName != "" {
|
||||
nickNameMap[contact.NickName] = contact
|
||||
}
|
||||
} else {
|
||||
chatRoomUserMap[contact.UserName] = contact
|
||||
}
|
||||
|
||||
}
|
||||
rows.Close()
|
||||
|
||||
c.Contact = contactMap
|
||||
c.ChatRoomUsers = chatRoomUserMap
|
||||
c.Alias2Contack = aliasMap
|
||||
c.Remark2Contack = remarkMap
|
||||
c.NickName2Contack = nickNameMap
|
||||
}
|
||||
|
||||
func (c *Contact) loadChatRoom() {
|
||||
|
||||
chatRoomMap := make(map[string]*model.ChatRoom)
|
||||
rows, err := c.db.Query("SELECT ChatRoomName, Reserved2, RoomData FROM ChatRoom")
|
||||
if err != nil {
|
||||
log.Errorf("查询群聊失败: %v", err)
|
||||
return
|
||||
}
|
||||
for rows.Next() {
|
||||
var chatRoom model.ChatRoomV3
|
||||
if err := rows.Scan(
|
||||
&chatRoom.ChatRoomName,
|
||||
&chatRoom.Reserved2,
|
||||
&chatRoom.RoomData,
|
||||
); err != nil {
|
||||
log.Printf("警告: 扫描群聊行失败: %v", err)
|
||||
continue
|
||||
}
|
||||
chatRoomMap[chatRoom.ChatRoomName] = chatRoom.Wrap()
|
||||
}
|
||||
rows.Close()
|
||||
c.ChatRoom = chatRoomMap
|
||||
}
|
||||
|
||||
func (c *Contact) loadSession() {
|
||||
|
||||
sessions := make([]*model.Session, 0)
|
||||
rows, err := c.db.Query("SELECT strUsrName, nOrder, strNickName, strContent, nTime FROM Session ORDER BY nOrder DESC")
|
||||
if err != nil {
|
||||
log.Errorf("查询群聊失败: %v", err)
|
||||
return
|
||||
}
|
||||
for rows.Next() {
|
||||
var sessionV3 model.SessionV3
|
||||
if err := rows.Scan(
|
||||
&sessionV3.StrUsrName,
|
||||
&sessionV3.NOrder,
|
||||
&sessionV3.StrNickName,
|
||||
&sessionV3.StrContent,
|
||||
&sessionV3.NTime,
|
||||
); err != nil {
|
||||
log.Printf("警告: 扫描历史会话失败: %v", err)
|
||||
continue
|
||||
}
|
||||
session := sessionV3.Wrap()
|
||||
sessions = append(sessions, session)
|
||||
|
||||
}
|
||||
rows.Close()
|
||||
c.Sessions = sessions
|
||||
}
|
||||
|
||||
func (c *Contact) ListContact() ([]*model.Contact, error) {
|
||||
contacts := make([]*model.Contact, 0, len(c.Contact))
|
||||
for _, contact := range c.Contact {
|
||||
contacts = append(contacts, contact)
|
||||
}
|
||||
return contacts, nil
|
||||
}
|
||||
|
||||
func (c *Contact) ListChatRoom() ([]*model.ChatRoom, error) {
|
||||
chatRooms := make([]*model.ChatRoom, 0, len(c.ChatRoom))
|
||||
for _, chatRoom := range c.ChatRoom {
|
||||
chatRooms = append(chatRooms, chatRoom)
|
||||
}
|
||||
return chatRooms, nil
|
||||
}
|
||||
|
||||
func (c *Contact) GetContact(key string) *model.Contact {
|
||||
if contact, ok := c.Contact[key]; ok {
|
||||
return contact
|
||||
}
|
||||
if contact, ok := c.Alias2Contack[key]; ok {
|
||||
return contact
|
||||
}
|
||||
if contact, ok := c.Remark2Contack[key]; ok {
|
||||
return contact
|
||||
}
|
||||
if contact, ok := c.NickName2Contack[key]; ok {
|
||||
return contact
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Contact) GetChatRoom(name string) *model.ChatRoom {
|
||||
if chatRoom, ok := c.ChatRoom[name]; ok {
|
||||
return chatRoom
|
||||
}
|
||||
|
||||
if contact := c.GetContact(name); contact != nil {
|
||||
if chatRoom, ok := c.ChatRoom[contact.UserName]; ok {
|
||||
return chatRoom
|
||||
} else {
|
||||
// 被删除的群聊，在 ChatRoom 记录中没有了，但是能找到 Contact，做下 Mock
|
||||
return &model.ChatRoom{
|
||||
Name: contact.UserName,
|
||||
Remark: contact.Remark,
|
||||
NickName: contact.NickName,
|
||||
Users: make([]model.ChatRoomUser, 0),
|
||||
User2DisplayName: make(map[string]string),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Contact) GetSession(limit int) []*model.Session {
|
||||
if limit <= 0 {
|
||||
limit = len(c.Sessions)
|
||||
}
|
||||
|
||||
if len(c.Sessions) < limit {
|
||||
limit = len(c.Sessions)
|
||||
}
|
||||
return c.Sessions[:limit]
|
||||
}
|
||||
|
||||
func (c *Contact) getFullContact(userName string) *model.Contact {
|
||||
if contact := c.GetContact(userName); contact != nil {
|
||||
return contact
|
||||
}
|
||||
if contact, ok := c.ChatRoomUsers[userName]; ok {
|
||||
return contact
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Contact) fillChatRoomInfo() {
|
||||
for i := range c.ChatRoom {
|
||||
if contact := c.GetContact(c.ChatRoom[i].Name); contact != nil {
|
||||
c.ChatRoom[i].Remark = contact.Remark
|
||||
c.ChatRoom[i].NickName = contact.NickName
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Contact) MessageFillInfo(msg *model.Message) {
|
||||
talker := msg.Talker
|
||||
if msg.IsChatRoom {
|
||||
talker = msg.ChatRoomSender
|
||||
if chatRoom := c.GetChatRoom(msg.Talker); chatRoom != nil {
|
||||
msg.CharRoomName = chatRoom.DisplayName()
|
||||
if displayName, ok := chatRoom.User2DisplayName[talker]; ok {
|
||||
msg.DisplayName = displayName
|
||||
}
|
||||
}
|
||||
}
|
||||
if msg.DisplayName == "" && msg.IsSender != 1 {
|
||||
if contact := c.getFullContact(talker); contact != nil {
|
||||
msg.DisplayName = contact.DisplayName()
|
||||
}
|
||||
}
|
||||
}
|
||||
321
internal/wechatdb/message.go
Normal file
321
internal/wechatdb/message.go
Normal file
@@ -0,0 +1,321 @@
|
||||
package wechatdb
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"log"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/sjzar/chatlog/pkg/model"
|
||||
"github.com/sjzar/chatlog/pkg/util"
|
||||
|
||||
_ "github.com/mattn/go-sqlite3"
|
||||
)
|
||||
|
||||
const (
|
||||
MessageFileV3 = "^MSG([0-9]?[0-9])?\\.db$"
|
||||
MessageFileV4 = "^messages_([0-9]?[0-9])+\\.db$"
|
||||
)
|
||||
|
||||
type Message struct {
|
||||
version int
|
||||
files []MsgDBInfo
|
||||
dbs map[string]*sql.DB
|
||||
}
|
||||
|
||||
type MsgDBInfo struct {
|
||||
FilePath string
|
||||
StartTime time.Time
|
||||
EndTime time.Time
|
||||
TalkerMap map[string]int
|
||||
}
|
||||
|
||||
func NewMessage(path string, version int) (*Message, error) {
|
||||
m := &Message{
|
||||
version: version,
|
||||
files: make([]MsgDBInfo, 0),
|
||||
dbs: make(map[string]*sql.DB),
|
||||
}
|
||||
|
||||
// 查找所有 MSG[0-13].db 文件
|
||||
files, err := util.FindFilesWithPatterns(path, MessageFileV3, true)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查找数据库文件失败: %v", err)
|
||||
}
|
||||
|
||||
if len(files) == 0 {
|
||||
return nil, fmt.Errorf("未找到任何数据库文件: %s", path)
|
||||
}
|
||||
|
||||
// 处理每个数据库文件
|
||||
for _, filePath := range files {
|
||||
// 连接数据库
|
||||
db, err := sql.Open("sqlite3", filePath)
|
||||
if err != nil {
|
||||
log.Printf("警告: 连接数据库 %s 失败: %v", filePath, err)
|
||||
continue
|
||||
}
|
||||
|
||||
// 获取 DBInfo 表中的开始时间
|
||||
// 首先检查表结构
|
||||
var startTime time.Time
|
||||
|
||||
// 尝试从 DBInfo 表中查找 Start Time 对应的记录
|
||||
rows, err := db.Query("SELECT tableIndex, tableVersion, tableDesc FROM DBInfo")
|
||||
if err != nil {
|
||||
log.Printf("警告: 查询数据库 %s 的 DBInfo 表失败: %v", filePath, err)
|
||||
db.Close()
|
||||
continue
|
||||
}
|
||||
|
||||
for rows.Next() {
|
||||
var tableIndex int
|
||||
var tableVersion int64
|
||||
var tableDesc string
|
||||
|
||||
if err := rows.Scan(&tableIndex, &tableVersion, &tableDesc); err != nil {
|
||||
log.Printf("警告: 扫描 DBInfo 行失败: %v", err)
|
||||
continue
|
||||
}
|
||||
|
||||
// 查找描述为 "Start Time" 的记录
|
||||
if strings.Contains(tableDesc, "Start Time") {
|
||||
startTime = time.Unix(tableVersion/1000, (tableVersion%1000)*1000000)
|
||||
break
|
||||
}
|
||||
}
|
||||
rows.Close()
|
||||
|
||||
// 组织 TalkerMap
|
||||
talkerMap := make(map[string]int)
|
||||
rows, err = db.Query("SELECT UsrName FROM Name2ID")
|
||||
if err != nil {
|
||||
log.Printf("警告: 查询数据库 %s 的 Name2ID 表失败: %v", filePath, err)
|
||||
db.Close()
|
||||
continue
|
||||
}
|
||||
|
||||
i := 1
|
||||
for rows.Next() {
|
||||
var userName string
|
||||
if err := rows.Scan(&userName); err != nil {
|
||||
log.Printf("警告: 扫描 Name2ID 行失败: %v", err)
|
||||
continue
|
||||
}
|
||||
talkerMap[userName] = i
|
||||
i++
|
||||
}
|
||||
|
||||
// 保存数据库信息
|
||||
m.files = append(m.files, MsgDBInfo{
|
||||
FilePath: filePath,
|
||||
StartTime: startTime,
|
||||
TalkerMap: talkerMap,
|
||||
})
|
||||
|
||||
// 保存数据库连接
|
||||
m.dbs[filePath] = db
|
||||
}
|
||||
|
||||
// 按照 StartTime 排序数据库文件
|
||||
sort.Slice(m.files, func(i, j int) bool {
|
||||
return m.files[i].StartTime.Before(m.files[j].StartTime)
|
||||
})
|
||||
|
||||
for i := range m.files {
|
||||
if i == len(m.files)-1 {
|
||||
m.files[i].EndTime = time.Now()
|
||||
} else {
|
||||
m.files[i].EndTime = m.files[i+1].StartTime
|
||||
}
|
||||
}
|
||||
|
||||
return m, nil
|
||||
}
|
||||
|
||||
// GetMessages 根据时间段和 talker 查询聊天记录
|
||||
func (m *Message) GetMessages(startTime, endTime time.Time, talker string, limit, offset int) ([]*model.Message, error) {
|
||||
// 找到时间范围内的数据库文件
|
||||
dbInfos := m.getDBInfosForTimeRange(startTime, endTime)
|
||||
if len(dbInfos) == 0 {
|
||||
return nil, fmt.Errorf("未找到时间范围 %v 到 %v 内的数据库文件", startTime, endTime)
|
||||
}
|
||||
|
||||
if len(dbInfos) == 1 {
|
||||
// LIMIT 和 OFFSET 逻辑在单文件情况下可以直接在 SQL 里处理
|
||||
return m.getMessagesSingleFile(dbInfos[0], startTime, endTime, talker, limit, offset)
|
||||
}
|
||||
|
||||
// 从每个相关数据库中查询消息
|
||||
totalMessages := []*model.Message{}
|
||||
|
||||
for _, dbInfo := range dbInfos {
|
||||
db, ok := m.dbs[dbInfo.FilePath]
|
||||
if !ok {
|
||||
log.Printf("警告: 数据库 %s 未打开", dbInfo.FilePath)
|
||||
continue
|
||||
}
|
||||
|
||||
// 构建查询条件
|
||||
// 使用 Sequence 查询，有索引
|
||||
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 Sequence, CreateTime, TalkerId, StrTalker, IsSender,
|
||||
Type, SubType, StrContent, CompressContent, BytesExtra
|
||||
FROM MSG
|
||||
WHERE %s
|
||||
ORDER BY Sequence ASC
|
||||
`, strings.Join(conditions, " AND "))
|
||||
|
||||
// 执行查询
|
||||
rows, err := db.Query(query, args...)
|
||||
if err != nil {
|
||||
log.Printf("警告: 查询数据库 %s 失败: %v", dbInfo.FilePath, err)
|
||||
continue
|
||||
}
|
||||
|
||||
// 处理查询结果
|
||||
for rows.Next() {
|
||||
var msg model.MessageV3
|
||||
var compressContent []byte
|
||||
var bytesExtra []byte
|
||||
|
||||
err := rows.Scan(
|
||||
&msg.Sequence,
|
||||
&msg.CreateTime,
|
||||
&msg.TalkerID,
|
||||
&msg.StrTalker,
|
||||
&msg.IsSender,
|
||||
&msg.Type,
|
||||
&msg.SubType,
|
||||
&msg.StrContent,
|
||||
&compressContent,
|
||||
&bytesExtra,
|
||||
)
|
||||
if err != nil {
|
||||
log.Printf("警告: 扫描消息行失败: %v", err)
|
||||
continue
|
||||
}
|
||||
msg.CompressContent = compressContent
|
||||
msg.BytesExtra = bytesExtra
|
||||
|
||||
totalMessages = append(totalMessages, msg.Wrap())
|
||||
}
|
||||
rows.Close()
|
||||
|
||||
if limit+offset > 0 && len(totalMessages) >= limit+offset {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
// 对所有消息按时间排序
|
||||
sort.Slice(totalMessages, func(i, j int) bool {
|
||||
return totalMessages[i].Sequence < totalMessages[j].Sequence
|
||||
})
|
||||
|
||||
// FIXME limit 和 offset 逻辑，在多文件边界条件下不好处理，直接查询全量数据后在进程里处理
|
||||
if limit > 0 {
|
||||
if offset >= len(totalMessages) {
|
||||
return []*model.Message{}, nil
|
||||
}
|
||||
end := offset + limit
|
||||
if end > len(totalMessages) || limit == 0 {
|
||||
end = len(totalMessages)
|
||||
}
|
||||
return totalMessages[offset:end], nil
|
||||
}
|
||||
|
||||
return totalMessages, nil
|
||||
|
||||
}
|
||||
|
||||
func (m *Message) getMessagesSingleFile(dbInfo MsgDBInfo, startTime, endTime time.Time, talker string, limit, offset int) ([]*model.Message, error) {
|
||||
// 构建查询条件
|
||||
// 使用 Sequence 查询，有索引
|
||||
conditions := []string{"Sequence >= ? AND Sequence <= ?"}
|
||||
args := []interface{}{startTime.Unix() * 1000, endTime.Unix() * 1000}
|
||||
if len(talker) > 0 {
|
||||
// TalkerId 有索引，优先使用
|
||||
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 Sequence, CreateTime, TalkerId, StrTalker, IsSender,
|
||||
Type, SubType, StrContent, CompressContent, BytesExtra
|
||||
FROM MSG
|
||||
WHERE %s
|
||||
ORDER BY Sequence ASC
|
||||
`, strings.Join(conditions, " AND "))
|
||||
|
||||
if limit > 0 {
|
||||
query += fmt.Sprintf(" LIMIT %d", limit)
|
||||
|
||||
if offset > 0 {
|
||||
query += fmt.Sprintf(" OFFSET %d", offset)
|
||||
}
|
||||
}
|
||||
|
||||
// 执行查询
|
||||
rows, err := m.dbs[dbInfo.FilePath].Query(query, args...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询数据库 %s 失败: %v", dbInfo.FilePath, err)
|
||||
}
|
||||
defer rows.Close()
|
||||
// 处理查询结果
|
||||
totalMessages := []*model.Message{}
|
||||
for rows.Next() {
|
||||
var msg model.MessageV3
|
||||
var compressContent []byte
|
||||
var bytesExtra []byte
|
||||
err := rows.Scan(
|
||||
&msg.Sequence,
|
||||
&msg.CreateTime,
|
||||
&msg.TalkerID,
|
||||
&msg.StrTalker,
|
||||
&msg.IsSender,
|
||||
&msg.Type,
|
||||
&msg.SubType,
|
||||
&msg.StrContent,
|
||||
&compressContent,
|
||||
&bytesExtra,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("扫描消息行失败: %v", err)
|
||||
}
|
||||
msg.CompressContent = compressContent
|
||||
msg.BytesExtra = bytesExtra
|
||||
totalMessages = append(totalMessages, msg.Wrap())
|
||||
}
|
||||
return totalMessages, nil
|
||||
}
|
||||
|
||||
func (m *Message) getDBInfosForTimeRange(startTime, endTime time.Time) []MsgDBInfo {
|
||||
var dbs []MsgDBInfo
|
||||
for _, info := range m.files {
|
||||
if info.StartTime.Before(endTime) && info.EndTime.After(startTime) {
|
||||
dbs = append(dbs, info)
|
||||
}
|
||||
}
|
||||
return dbs
|
||||
}
|
||||
117
internal/wechatdb/wechatdb.go
Normal file
117
internal/wechatdb/wechatdb.go
Normal file
@@ -0,0 +1,117 @@
|
||||
package wechatdb
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/sjzar/chatlog/pkg/model"
|
||||
|
||||
_ "github.com/mattn/go-sqlite3"
|
||||
)
|
||||
|
||||
type DB struct {
|
||||
BasePath string
|
||||
Version int
|
||||
|
||||
contact *Contact
|
||||
message *Message
|
||||
}
|
||||
|
||||
func New(path string, version int) (*DB, error) {
|
||||
w := &DB{
|
||||
BasePath: path,
|
||||
Version: version,
|
||||
}
|
||||
|
||||
// 初始化，加载数据库文件信息
|
||||
if err := w.Initialize(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return w, nil
|
||||
}
|
||||
|
||||
func (w *DB) Close() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *DB) Initialize() error {
|
||||
|
||||
var err error
|
||||
w.message, err = NewMessage(w.BasePath, w.Version)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
w.contact, err = NewContact(w.BasePath, w.Version)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *DB) GetMessages(start, end time.Time, talker string, limit, offset int) ([]*model.Message, error) {
|
||||
|
||||
if talker != "" {
|
||||
if contact := w.contact.GetContact(talker); contact != nil {
|
||||
talker = contact.UserName
|
||||
}
|
||||
}
|
||||
|
||||
messages, err := w.message.GetMessages(start, end, talker, limit, offset)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for i := range messages {
|
||||
w.contact.MessageFillInfo(messages[i])
|
||||
}
|
||||
|
||||
return messages, nil
|
||||
}
|
||||
|
||||
type ListContactResp struct {
|
||||
Items []*model.Contact `json:"items"`
|
||||
}
|
||||
|
||||
func (w *DB) ListContact() (*ListContactResp, error) {
|
||||
list, err := w.contact.ListContact()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &ListContactResp{
|
||||
Items: list,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (w *DB) GetContact(userName string) *model.Contact {
|
||||
return w.contact.GetContact(userName)
|
||||
}
|
||||
|
||||
type ListChatRoomResp struct {
|
||||
Items []*model.ChatRoom `json:"items"`
|
||||
}
|
||||
|
||||
func (w *DB) ListChatRoom() (*ListChatRoomResp, error) {
|
||||
list, err := w.contact.ListChatRoom()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &ListChatRoomResp{
|
||||
Items: list,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (w *DB) GetChatRoom(userName string) *model.ChatRoom {
|
||||
return w.contact.GetChatRoom(userName)
|
||||
}
|
||||
|
||||
type GetSessionResp struct {
|
||||
Items []*model.Session `json:"items"`
|
||||
}
|
||||
|
||||
func (w *DB) GetSession(limit int) (*GetSessionResp, error) {
|
||||
sessions := w.contact.GetSession(limit)
|
||||
return &GetSessionResp{
|
||||
Items: sessions,
|
||||
}, nil
|
||||
}
|
||||
Reference in New Issue
Block a user