agents_and_robots/shell/matrix/client.go

// Package matrix wraps mautrix-go for agent use.
package matrix

import (
	"context"
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"log/slog"
	"os"
	"path/filepath"
	"strings"

	"github.com/yuin/goldmark"
	"github.com/yuin/goldmark/extension"
	"maunium.net/go/mautrix"
	"maunium.net/go/mautrix/crypto"
	"maunium.net/go/mautrix/crypto/cryptohelper"
	"maunium.net/go/mautrix/crypto/ssss"
	"maunium.net/go/mautrix/event"
	"maunium.net/go/mautrix/id"

	"github.com/enmanuel/agents/internal/config"
)

// Client wraps a mautrix client with agent-relevant helpers.
type Client struct {
	raw *mautrix.Client
	cfg config.MatrixCfg
}

// New creates and authenticates a Matrix client from config.
// The access token is read from the env var specified in cfg.AccessTokenEnv.
func New(cfg config.MatrixCfg) (*Client, error) {
	token := os.Getenv(cfg.AccessTokenEnv)
	if token == "" {
		return nil, fmt.Errorf("env var %s is not set", cfg.AccessTokenEnv)
	}

	raw, err := mautrix.NewClient(cfg.Homeserver, id.UserID(cfg.UserID), token)
	if err != nil {
		return nil, fmt.Errorf("create matrix client: %w", err)
	}

	if cfg.DeviceID != "" {
		raw.DeviceID = id.DeviceID(cfg.DeviceID)
	}

	return &Client{raw: raw, cfg: cfg}, nil
}

// cryptoIniter abstracts crypto helper creation for testing.
type cryptoIniter interface {
	newHelper(pickleKey []byte, storePath string) (cryptoHelper, error)
}

// cryptoHelper abstracts the mautrix CryptoHelper for testing.
type cryptoHelper interface {
	io.Closer
	Init(ctx context.Context) error
	SetAccountID(id string)
}

// mautrixCryptoIniter is the real implementation using mautrix.
type mautrixCryptoIniter struct {
	raw *mautrix.Client
}

func (m *mautrixCryptoIniter) newHelper(pickleKey []byte, storePath string) (cryptoHelper, error) {
	h, err := cryptohelper.NewCryptoHelper(m.raw, pickleKey, storePath)
	if err != nil {
		return nil, err
	}
	return &mautrixCryptoWrapper{h}, nil
}

type mautrixCryptoWrapper struct {
	*cryptohelper.CryptoHelper
}

func (w *mautrixCryptoWrapper) SetAccountID(id string) { w.DBAccountID = id }

// InitCrypto sets up end-to-end encryption using the mautrix cryptohelper.
// storePath is the SQLite file path for crypto material (e.g. "./agents/<id>/data/crypto/crypto.db").
// pickleKeyHex is a hex-encoded key for encrypting crypto material at rest. If empty,
// falls back to sha256(access_token) for backward compatibility.
// agentID namespaces the crypto state within the database.
// Returns an io.Closer that must be called on agent shutdown to flush the crypto store.
func (c *Client) InitCrypto(ctx context.Context, storePath, pickleKeyHex, agentID string) (io.Closer, error) {
	// Resolve the actual device ID from the server.
	whoami, err := c.raw.Whoami(ctx)
	if err != nil {
		return nil, fmt.Errorf("whoami for crypto init: %w", err)
	}
	c.raw.DeviceID = whoami.DeviceID

	initer := &mautrixCryptoIniter{raw: c.raw}
	closer, helper, err := initCryptoCore(ctx, storePath, pickleKeyHex, c.raw.AccessToken, agentID, initer, slog.Default())
	if err != nil {
		return nil, err
	}

	// Assign the real mautrix crypto helper — this satisfies mautrix.CryptoHelper.
	wrapper := helper.(*mautrixCryptoWrapper)
	c.raw.Crypto = wrapper

	// Log E2EE diagnostic state to help debug verification issues.
	logCryptoDiagnostics(ctx, wrapper, c.raw, slog.Default())

	return closer, nil
}

// ssssKeyFetcher abstracts the SSSS + cross-signing key retrieval for testing.
type ssssKeyFetcher interface {
	GetDefaultKeyData(ctx context.Context) (string, ssssKeyVerifier, error)
	FetchCrossSigningKeysFromSSSS(ctx context.Context, key *ssss.Key) error
}

// ssssKeyVerifier abstracts the SSSS key metadata verification.
type ssssKeyVerifier interface {
	VerifyRecoveryKey(keyID, recoveryKey string) (*ssss.Key, error)
}

// olmSSSSFetcher adapts *crypto.OlmMachine to the ssssKeyFetcher interface.
type olmSSSSFetcher struct {
	machine *crypto.OlmMachine
}

func (o *olmSSSSFetcher) GetDefaultKeyData(ctx context.Context) (string, ssssKeyVerifier, error) {
	keyID, keyData, err := o.machine.SSSS.GetDefaultKeyData(ctx)
	return keyID, keyData, err
}

func (o *olmSSSSFetcher) FetchCrossSigningKeysFromSSSS(ctx context.Context, key *ssss.Key) error {
	return o.machine.FetchCrossSigningKeysFromSSSS(ctx, key)
}

// FetchCrossSigningKeys retrieves cross-signing private keys from SSSS
// (server-side secret storage) using the given base58 recovery key.
// This allows the agent to sign its own device, eliminating the
// "Encrypted by a device not verified by its owner" warning.
func (c *Client) FetchCrossSigningKeys(ctx context.Context, recoveryKey string) error {
	wrapper, ok := c.raw.Crypto.(*mautrixCryptoWrapper)
	if !ok || wrapper == nil {
		return fmt.Errorf("crypto not initialized")
	}
	machine := wrapper.Machine()
	if machine == nil {
		return fmt.Errorf("olm machine not available")
	}
	return fetchCrossSigningKeysCore(ctx, &olmSSSSFetcher{machine}, recoveryKey)
}

// SignOwnDevice signs the bot's current device with the self-signing key.
// This is the step that makes Element show the device as "verified".
// Must be called after cross-signing private keys are available (via
// FetchCrossSigningKeys or GenerateAndUploadCrossSigningKeys).
// It force-fetches device keys from the server first to ensure the local
// store has the correct signing key.
func (c *Client) SignOwnDevice(ctx context.Context) error {
	wrapper, ok := c.raw.Crypto.(*mautrixCryptoWrapper)
	if !ok || wrapper == nil {
		return fmt.Errorf("crypto not initialized")
	}
	machine := wrapper.Machine()
	if machine == nil {
		return fmt.Errorf("olm machine not available")
	}

	// Force-fetch own device keys so the local store has the correct signing key.
	// Without this, SignOwnDevice fails with "different signing key" when the
	// store has a stale or empty entry.
	devices, err := machine.FetchKeys(ctx, []id.UserID{c.raw.UserID}, true)
	if err != nil {
		return fmt.Errorf("fetch own device keys: %w", err)
	}
	userDevices, ok := devices[c.raw.UserID]
	if !ok {
		return fmt.Errorf("own user not found in fetched keys")
	}
	device, ok := userDevices[c.raw.DeviceID]
	if !ok {
		return fmt.Errorf("own device %s not found in fetched keys", c.raw.DeviceID)
	}

	return machine.SignOwnDevice(ctx, device)
}

// fetchCrossSigningKeysCore contains the testable logic for SSSS key retrieval.
func fetchCrossSigningKeysCore(ctx context.Context, fetcher ssssKeyFetcher, recoveryKey string) error {
	keyID, keyData, err := fetcher.GetDefaultKeyData(ctx)
	if err != nil {
		return fmt.Errorf("get SSSS default key: %w", err)
	}

	key, err := keyData.VerifyRecoveryKey(keyID, recoveryKey)
	if err != nil {
		return fmt.Errorf("verify recovery key: %w", err)
	}

	if err := fetcher.FetchCrossSigningKeysFromSSSS(ctx, key); err != nil {
		return fmt.Errorf("fetch cross-signing keys from SSSS: %w", err)
	}

	return nil
}

// initCryptoCore contains the testable logic: pickle key resolution, store
// creation, and auto-recovery on stale crypto.db. Returns (closer, helper, err).
func initCryptoCore(ctx context.Context, storePath, pickleKeyHex, accessToken, agentID string, initer cryptoIniter, logger *slog.Logger) (io.Closer, cryptoHelper, error) {
	pickleKey, err := resolvePickleKey(pickleKeyHex, accessToken)
	if err != nil {
		return nil, nil, err
	}

	if err := os.MkdirAll(filepath.Dir(storePath), 0700); err != nil {
		return nil, nil, fmt.Errorf("create crypto store dir: %w", err)
	}

	helper, err := initHelper(ctx, initer, pickleKey, storePath, agentID)
	if err != nil && strings.Contains(err.Error(), "not marked as shared") {
		logger.Warn("crypto store inconsistent, attempting auto-recovery",
			"store", storePath,
		)
		if removeErr := os.Remove(storePath); removeErr != nil && !os.IsNotExist(removeErr) {
			return nil, nil, fmt.Errorf("auto-recovery: remove stale crypto.db: %w (original: %w)", removeErr, err)
		}
		helper, err = initHelper(ctx, initer, pickleKey, storePath, agentID)
		if err != nil {
			return nil, nil, fmt.Errorf("e2ee init after auto-recovery: %w", err)
		}
		logger.Info("e2ee auto-recovery succeeded")
	} else if err != nil {
		return nil, nil, fmt.Errorf("init e2ee: %w", err)
	}

	return helper, helper, nil
}

func initHelper(ctx context.Context, initer cryptoIniter, pickleKey []byte, storePath, agentID string) (cryptoHelper, error) {
	helper, err := initer.newHelper(pickleKey, storePath)
	if err != nil {
		return nil, fmt.Errorf("create crypto helper: %w", err)
	}
	helper.SetAccountID(agentID)
	if err := helper.Init(ctx); err != nil {
		return nil, err
	}
	return helper, nil
}

// resolvePickleKey decodes a hex key or derives one from the access token.
func resolvePickleKey(pickleKeyHex, accessToken string) ([]byte, error) {
	if pickleKeyHex != "" {
		key, err := hex.DecodeString(pickleKeyHex)
		if err != nil {
			return nil, fmt.Errorf("decode pickle_key_env: %w", err)
		}
		return key, nil
	}
	sum := sha256.Sum256([]byte(accessToken))
	return sum[:], nil
}

// SendText sends a plain-text message to a room.
// If the room has E2EE enabled and crypto is initialized, the message is encrypted automatically.
func (c *Client) SendText(ctx context.Context, roomID, text string) error {
	_, err := c.raw.SendText(ctx, id.RoomID(roomID), text)
	return err
}

// SendMarkdown sends a formatted (Markdown) message to a room.
// Body contains the raw markdown (plaintext fallback per Matrix spec).
// FormattedBody contains the HTML rendered by goldmark.
func (c *Client) SendMarkdown(ctx context.Context, roomID, markdown string) error {
	html := mdToHTML(markdown)
	content := &event.MessageEventContent{
		MsgType:       event.MsgText,
		Body:          markdown,
		Format:        event.FormatHTML,
		FormattedBody: html,
	}
	_, err := c.raw.SendMessageEvent(ctx, id.RoomID(roomID), event.EventMessage, content)
	return err
}

// SendMarkdownGetID sends a formatted (Markdown) message to a room and returns
// the event ID of the sent message. Useful for later editing via EditMessage.
func (c *Client) SendMarkdownGetID(ctx context.Context, roomID, markdown string) (string, error) {
	html := mdToHTML(markdown)
	content := &event.MessageEventContent{
		MsgType:       event.MsgText,
		Body:          markdown,
		Format:        event.FormatHTML,
		FormattedBody: html,
	}
	resp, err := c.raw.SendMessageEvent(ctx, id.RoomID(roomID), event.EventMessage, content)
	if err != nil {
		return "", err
	}
	return resp.EventID.String(), nil
}

// EditMessage edits a previously sent message in a room using m.replace relation.
// originalEventID is the event ID of the message to replace.
// The new content is rendered from markdown.
func (c *Client) EditMessage(ctx context.Context, roomID, originalEventID, markdown string) error {
	html := mdToHTML(markdown)

	// Matrix spec: m.new_content holds the replacement, m.relates_to with
	// rel_type=m.replace points to the original event.
	content := &event.MessageEventContent{
		MsgType:       event.MsgText,
		Body:          "* " + markdown, // per spec: prefix with "* " for fallback
		Format:        event.FormatHTML,
		FormattedBody: "* " + html,
		RelatesTo: &event.RelatesTo{
			Type:    event.RelReplace,
			EventID: id.EventID(originalEventID),
		},
		NewContent: &event.MessageEventContent{
			MsgType:       event.MsgText,
			Body:          markdown,
			Format:        event.FormatHTML,
			FormattedBody: html,
		},
	}
	_, err := c.raw.SendMessageEvent(ctx, id.RoomID(roomID), event.EventMessage, content)
	return err
}

// mdToHTML converts a Markdown string to HTML using goldmark with full extensions.
var mdParser = goldmark.New(
	goldmark.WithExtensions(
		extension.GFM,
		extension.DefinitionList,
		extension.Footnote,
		extension.Typographer,
		extension.CJK,
	),
)

func mdToHTML(md string) string {
	var buf bytes.Buffer
	if err := mdParser.Convert([]byte(md), &buf); err != nil {
		return md // fallback to raw markdown on error
	}
	return buf.String()
}

// SendReplyMarkdown sends a formatted message as a reply to a specific event.
// It sets m.in_reply_to so Matrix clients show the original message as a quote.
func (c *Client) SendReplyMarkdown(ctx context.Context, roomID, inReplyTo, markdown string) error {
	html := mdToHTML(markdown)
	content := &event.MessageEventContent{
		MsgType:       event.MsgText,
		Body:          markdown,
		Format:        event.FormatHTML,
		FormattedBody: html,
		RelatesTo:     (&event.RelatesTo{}).SetReplyTo(id.EventID(inReplyTo)),
	}
	_, err := c.raw.SendMessageEvent(ctx, id.RoomID(roomID), event.EventMessage, content)
	return err
}

// SendThreadMarkdown sends a formatted message as part of a Matrix thread.
// threadRootID is the event that started the thread (always the same for all messages in a thread).
// inReplyTo is the specific event being replied to within the thread (used as fallback for non-thread clients).
// If inReplyTo is empty, it defaults to threadRootID.
func (c *Client) SendThreadMarkdown(ctx context.Context, roomID, threadRootID, inReplyTo, markdown string) error {
	if inReplyTo == "" {
		inReplyTo = threadRootID
	}
	html := mdToHTML(markdown)
	// Must use a pointer so mautrix-go can call OptionalGetRelatesTo() (pointer receiver)
	// and copy m.relates_to to the outer m.room.encrypted event. Without this, Element
	// Web does not see the thread relationship and shows the reply in the main timeline.
	content := &event.MessageEventContent{
		MsgType:       event.MsgText,
		Body:          markdown,
		Format:        event.FormatHTML,
		FormattedBody: html,
		RelatesTo:     (&event.RelatesTo{}).SetThread(id.EventID(threadRootID), id.EventID(inReplyTo)),
	}
	_, err := c.raw.SendMessageEvent(ctx, id.RoomID(roomID), event.EventMessage, content)
	return err
}

// SendReaction sends a reaction to an event.
func (c *Client) SendReaction(ctx context.Context, roomID, eventID, reaction string) error {
	_, err := c.raw.SendReaction(ctx, id.RoomID(roomID), id.EventID(eventID), reaction)
	return err
}

// SendTyping sets the typing indicator in a room.
func (c *Client) SendTyping(ctx context.Context, roomID string, typing bool) error {
	_, err := c.raw.UserTyping(ctx, id.RoomID(roomID), typing, 5000)
	return err
}

// diagMachine abstracts the crypto.OlmMachine methods used by diagnostics,
// allowing unit tests to substitute a fake without a real crypto store.
type diagMachine interface {
	GetOwnCrossSigningPublicKeys(ctx context.Context) *crypto.CrossSigningPublicKeysCache
	OwnIdentity() *id.Device
	ExportCrossSigningKeys() crypto.CrossSigningSeeds
	ResolveTrustContext(ctx context.Context, device *id.Device) (id.TrustState, error)
	IsDeviceTrusted(device *id.Device) bool
}

// logCryptoDiagnostics logs the E2EE state after initialization.
// This helps diagnose "Encrypted by a device not verified by its owner" warnings.
// It is defensive against nil fields in the crypto machine to avoid panics.
func logCryptoDiagnostics(ctx context.Context, wrapper *mautrixCryptoWrapper, raw *mautrix.Client, logger *slog.Logger) {
	machine := wrapper.Machine()
	if machine == nil {
		logger.Warn("e2ee diagnostics: olm machine is nil")
		return
	}
	logCryptoDiagnosticsCore(ctx, machine, raw.UserID, raw.DeviceID, logger)
}

// logCryptoDiagnosticsCore contains the testable diagnostics logic.
func logCryptoDiagnosticsCore(ctx context.Context, machine diagMachine, userID id.UserID, deviceID id.DeviceID, logger *slog.Logger) {
	logger.Info("e2ee diagnostics: device info",
		"user_id", userID,
		"device_id", deviceID,
	)

	// Check own cross-signing public keys
	ownKeys := machine.GetOwnCrossSigningPublicKeys(ctx)
	if ownKeys == nil {
		logger.Warn("e2ee diagnostics: NO cross-signing public keys found — this is likely why messages show 'not verified by its owner'",
			"hint", "run: go run -tags goolm ./cmd/verify --homeserver <hs> --username <user> --password <pass> --token <token>",
		)
	} else {
		logger.Info("e2ee diagnostics: cross-signing public keys found",
			"master_key", truncateKey(string(ownKeys.MasterKey)),
			"self_signing_key", truncateKey(string(ownKeys.SelfSigningKey)),
			"user_signing_key", truncateKey(string(ownKeys.UserSigningKey)),
		)
	}

	// Check if our own device is trusted via cross-signing
	ownDevice := machine.OwnIdentity()
	if ownDevice == nil {
		logger.Warn("e2ee diagnostics: own device identity is nil — cannot check trust state")
	} else {
		logDeviceTrust(ctx, machine, ownDevice, logger)
	}

	// Check if cross-signing private keys are available (needed to sign devices).
	// ExportCrossSigningKeys panics when CrossSigningKeys is nil (no private keys
	// loaded), so we guard the call.
	logCrossSigningSeeds(machine, logger)
}

// logCrossSigningSeeds safely exports and logs cross-signing private key availability.
// mautrix's ExportCrossSigningKeys panics when CrossSigningKeys is nil, so we
// recover from the panic instead of relying on unexported internal fields.
func logCrossSigningSeeds(machine diagMachine, logger *slog.Logger) {
	var (
		hasMasterSeed      bool
		hasSelfSigningSeed bool
		hasUserSigningSeed bool
		recovered          bool
	)

	func() {
		defer func() {
			if r := recover(); r != nil {
				recovered = true
			}
		}()
		seeds := machine.ExportCrossSigningKeys()
		hasMasterSeed = len(seeds.MasterKey) > 0
		hasSelfSigningSeed = len(seeds.SelfSigningKey) > 0
		hasUserSigningSeed = len(seeds.UserSigningKey) > 0
	}()

	if recovered {
		logger.Warn("e2ee diagnostics: cross-signing private keys not available (not loaded in crypto store)")
		return
	}

	logger.Info("e2ee diagnostics: cross-signing private keys in store",
		"master_seed", hasMasterSeed,
		"self_signing_seed", hasSelfSigningSeed,
		"user_signing_seed", hasUserSigningSeed,
	)

	if !hasSelfSigningSeed {
		logger.Warn("e2ee diagnostics: self-signing private key NOT in store — the bot cannot sign its own device",
			"hint", "run cmd/verify with the SAME crypto store path the agent uses",
		)
	}
}

// logDeviceTrust resolves and logs the trust state for a device.
func logDeviceTrust(ctx context.Context, machine diagMachine, device *id.Device, logger *slog.Logger) {
	trust, err := machine.ResolveTrustContext(ctx, device)
	if err != nil {
		logger.Warn("e2ee diagnostics: failed to resolve device trust",
			"device_id", device.DeviceID,
			"err", err,
		)
		return
	}

	logger.Info("e2ee diagnostics: own device trust state",
		"device_id", device.DeviceID,
		"trust_state", trust.String(),
		"is_trusted", machine.IsDeviceTrusted(device),
	)

	if trust < id.TrustStateCrossSignedTOFU {
		logger.Warn("e2ee diagnostics: device is NOT cross-signed — recipients will see 'not verified by its owner'",
			"trust_state", trust.String(),
			"required_minimum", "CrossSignedTOFU",
		)
	}
}

// truncateKey returns first 8 chars of a key for safe logging.
func truncateKey(key string) string {
	if len(key) > 8 {
		return key[:8] + "..."
	}
	return key
}

// SetPresence sets the bot's presence status (online, unavailable, offline).
func (c *Client) SetPresence(ctx context.Context, status event.Presence) error {
	return c.raw.SetPresence(ctx, status)
}

// Raw returns the underlying mautrix.Client for advanced use.
func (c *Client) Raw() *mautrix.Client {
	return c.raw
}