feat: load packet logs from configured source
This commit is contained in:
zhaoyilun
@@ -1,6 +1,10 @@
|
||||
package mcpserver
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
|
||||
"isphere-ai-bridge/internal/helperclient"
|
||||
@@ -9,18 +13,36 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
ServerName = "isphere-ai-bridge"
|
||||
ServerTitle = "iSphere AI Bridge"
|
||||
ServerVersion = "0.1.0"
|
||||
ServerName = "isphere-ai-bridge"
|
||||
ServerTitle = "iSphere AI Bridge"
|
||||
ServerVersion = "0.1.0"
|
||||
EnvPacketLogFileKey = "ISPHERE_PACKET_LOG_FILE"
|
||||
)
|
||||
|
||||
func NewServer() *mcp.Server {
|
||||
return NewServerWithReceiveSource(isphere.EncryptedPacketLogSource{})
|
||||
}
|
||||
|
||||
func NewServerFromEnv() (*mcp.Server, error) {
|
||||
var source tools.ReceiveMessagesSource = isphere.EncryptedPacketLogSource{}
|
||||
packetLogFile := strings.TrimSpace(os.Getenv(EnvPacketLogFileKey))
|
||||
if packetLogFile != "" {
|
||||
loaded, err := isphere.LoadEncryptedPacketLogSourceFromFile(packetLogFile)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("configure %s: %w", EnvPacketLogFileKey, err)
|
||||
}
|
||||
source = loaded
|
||||
}
|
||||
return NewServerWithReceiveSource(source), nil
|
||||
}
|
||||
|
||||
func NewServerWithReceiveSource(source tools.ReceiveMessagesSource) *mcp.Server {
|
||||
server := mcp.NewServer(&mcp.Implementation{
|
||||
Name: ServerName,
|
||||
Title: ServerTitle,
|
||||
Version: ServerVersion,
|
||||
}, nil)
|
||||
tools.RegisterWinHelperTools(server, helperclient.Client{})
|
||||
tools.RegisterISphereReadTools(server, isphere.EncryptedPacketLogSource{})
|
||||
tools.RegisterISphereReadTools(server, source)
|
||||
return server
|
||||
}
|
||||
|
||||
@@ -2,6 +2,11 @@ package mcpserver
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/cipher"
|
||||
"crypto/des"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"reflect"
|
||||
"sort"
|
||||
"testing"
|
||||
@@ -57,3 +62,127 @@ func TestNewServerRegistersN8WinHelperToolsWithoutCallingHelper(t *testing.T) {
|
||||
t.Fatalf("server tool names = %#v, want %#v", gotNames, wantNames)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewServerFromEnvUsesPacketLogFile(t *testing.T) {
|
||||
plaintext := `--------------------------------------------------------------------------------------------------------------------------------------------
|
||||
2026/7/7 15:30:07
|
||||
<message id="msg-env-1" from="sender@imopenfire1-lanzhou/imp_pc_4.1.2.6842" to="receiver@imopenfire1-lanzhou" type="chat">
|
||||
<body>redacted from env</body>
|
||||
<received xmlns="urn:xmpp:receipts" id="receipt-env-1" type="1" stamp="" />
|
||||
<subject>ENV_SOURCE</subject>
|
||||
<isphere xmlns="isphere.im" type="1001" sendtime="1783423807000" version="1" />
|
||||
<readed />
|
||||
</message>`
|
||||
filePath := writeEncryptedPacketLogFileForServerTest(t, plaintext)
|
||||
t.Setenv("ISPHERE_PACKET_LOG_FILE", filePath)
|
||||
|
||||
server, err := NewServerFromEnv()
|
||||
if err != nil {
|
||||
t.Fatalf("NewServerFromEnv returned error: %v", err)
|
||||
}
|
||||
session, cleanup := connectServerTestSession(t, server)
|
||||
defer cleanup()
|
||||
|
||||
callResult, err := session.CallTool(context.Background(), &mcp.CallToolParams{
|
||||
Name: "isphere_receive_messages",
|
||||
Arguments: map[string]any{"limit": 10},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("call isphere_receive_messages: %v", err)
|
||||
}
|
||||
if callResult.IsError {
|
||||
t.Fatalf("call result is error: %+v", callResult)
|
||||
}
|
||||
|
||||
var decoded struct {
|
||||
Messages []struct {
|
||||
ID string `json:"id"`
|
||||
Text string `json:"text"`
|
||||
Subject string `json:"subject"`
|
||||
ConversationID string `json:"conversation_id"`
|
||||
SenderID string `json:"sender_id"`
|
||||
ReceiverID string `json:"receiver_id"`
|
||||
Read bool `json:"read"`
|
||||
} `json:"messages"`
|
||||
}
|
||||
payload, err := json.Marshal(callResult.StructuredContent)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal structured content: %v", err)
|
||||
}
|
||||
if err := json.Unmarshal(payload, &decoded); err != nil {
|
||||
t.Fatalf("decode structured content %s: %v", payload, err)
|
||||
}
|
||||
if len(decoded.Messages) != 1 {
|
||||
t.Fatalf("message count = %d, want 1", len(decoded.Messages))
|
||||
}
|
||||
msg := decoded.Messages[0]
|
||||
if msg.ID != "msg-env-1" || msg.Text != "redacted from env" || msg.Subject != "ENV_SOURCE" {
|
||||
t.Fatalf("unexpected message: %+v", msg)
|
||||
}
|
||||
if msg.ConversationID != "sender@imopenfire1-lanzhou|receiver@imopenfire1-lanzhou" {
|
||||
t.Fatalf("conversation id = %q", msg.ConversationID)
|
||||
}
|
||||
if msg.SenderID != "sender@imopenfire1-lanzhou" || msg.ReceiverID != "receiver@imopenfire1-lanzhou" || !msg.Read {
|
||||
t.Fatalf("metadata = %+v", msg)
|
||||
}
|
||||
}
|
||||
|
||||
func connectServerTestSession(t *testing.T, server *mcp.Server) (*mcp.ClientSession, func()) {
|
||||
t.Helper()
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
serverTransport, clientTransport := mcp.NewInMemoryTransports()
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
errCh <- server.Run(ctx, serverTransport)
|
||||
}()
|
||||
|
||||
client := mcp.NewClient(&mcp.Implementation{Name: "mcpserver-env-test-client", Version: "0.0.0"}, nil)
|
||||
session, err := client.Connect(ctx, clientTransport, nil)
|
||||
if err != nil {
|
||||
cancel()
|
||||
t.Fatalf("connect client: %v", err)
|
||||
}
|
||||
|
||||
cleanup := func() {
|
||||
_ = session.Close()
|
||||
cancel()
|
||||
select {
|
||||
case <-errCh:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatalf("server did not stop")
|
||||
}
|
||||
}
|
||||
return session, cleanup
|
||||
}
|
||||
|
||||
func writeEncryptedPacketLogFileForServerTest(t *testing.T, plaintext string) string {
|
||||
t.Helper()
|
||||
path := t.TempDir() + "\\packet-reader.log"
|
||||
line := encryptPacketLogLineForServerTest(t, plaintext)
|
||||
if err := os.WriteFile(path, []byte("\n"+line+"\n"), 0o600); err != nil {
|
||||
t.Fatalf("write encrypted packet log file: %v", err)
|
||||
}
|
||||
return path
|
||||
}
|
||||
|
||||
func encryptPacketLogLineForServerTest(t *testing.T, plaintext string) string {
|
||||
t.Helper()
|
||||
block, err := des.NewCipher([]byte("hyhccdtm"))
|
||||
if err != nil {
|
||||
t.Fatalf("NewCipher: %v", err)
|
||||
}
|
||||
plain := padPacketLogLineForServerTest([]byte(plaintext), des.BlockSize)
|
||||
ciphertext := make([]byte, len(plain))
|
||||
iv := []byte{0x12, 0x34, 0x56, 0x78, 0x90, 0xAB, 0xCD, 0xEF}
|
||||
cipher.NewCBCEncrypter(block, iv).CryptBlocks(ciphertext, plain)
|
||||
return base64.StdEncoding.EncodeToString(ciphertext)
|
||||
}
|
||||
|
||||
func padPacketLogLineForServerTest(data []byte, blockSize int) []byte {
|
||||
pad := blockSize - len(data)%blockSize
|
||||
out := append([]byte(nil), data...)
|
||||
for i := 0; i < pad; i++ {
|
||||
out = append(out, byte(pad))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user