motr-enclave/main.go

package main

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"strings"

	"enclave/internal/keybase"
	"enclave/internal/state"
	"enclave/internal/types"

	"github.com/extism/go-pdk"
)

func main() { state.Default() }

//go:wasmexport ping
func ping() int32 {
	pdk.Log(pdk.LogInfo, "ping: received request")

	var input types.PingInput
	if err := pdk.InputJSON(&input); err != nil {
		output := types.PingOutput{
			Success: false,
			Message: fmt.Sprintf("failed to parse input: %s", err),
		}
		pdk.OutputJSON(output)
		return 0
	}

	output := types.PingOutput{
		Success: true,
		Message: "pong",
		Echo:    input.Message,
	}

	if err := pdk.OutputJSON(output); err != nil {
		pdk.Log(pdk.LogError, fmt.Sprintf("ping: failed to output: %s", err))
		return 1
	}

	pdk.Log(pdk.LogInfo, fmt.Sprintf("ping: responded with echo=%s", input.Message))
	return 0
}

//go:wasmexport generate
func generate() int32 {
	pdk.Log(pdk.LogInfo, "generate: starting database initialization")

	var input types.GenerateInput
	if err := pdk.InputJSON(&input); err != nil {
		pdk.SetError(fmt.Errorf("generate: failed to parse input: %w", err))
		return 1
	}

	if input.Credential == "" {
		pdk.SetError(errors.New("generate: credential is required"))
		return 1
	}

	credentialBytes, err := base64.StdEncoding.DecodeString(input.Credential)
	if err != nil {
		pdk.SetError(fmt.Errorf("generate: invalid base64 credential: %w", err))
		return 1
	}

	did, err := initializeDatabase(credentialBytes)
	if err != nil {
		pdk.SetError(fmt.Errorf("generate: failed to initialize database: %w", err))
		return 1
	}

	state.SetInitialized(true)
	state.SetDID(did)

	dbBytes, err := serializeDatabase()
	if err != nil {
		pdk.SetError(fmt.Errorf("generate: failed to serialize database: %w", err))
		return 1
	}

	output := types.GenerateOutput{
		DID:      did,
		Database: dbBytes,
	}

	if err := pdk.OutputJSON(output); err != nil {
		pdk.SetError(fmt.Errorf("generate: failed to output result: %w", err))
		return 1
	}

	pdk.Log(pdk.LogInfo, fmt.Sprintf("generate: created DID %s", did))
	return 0
}

//go:wasmexport load
func load() int32 {
	pdk.Log(pdk.LogInfo, "load: loading database from buffer")

	var input types.LoadInput
	if err := pdk.InputJSON(&input); err != nil {
		pdk.SetError(fmt.Errorf("load: failed to parse input: %w", err))
		return 1
	}

	if len(input.Database) == 0 {
		pdk.SetError(errors.New("load: database buffer is required"))
		return 1
	}

	did, err := loadDatabase(input.Database)
	if err != nil {
		output := types.LoadOutput{
			Success: false,
			Error:   err.Error(),
		}
		pdk.OutputJSON(output)
		return 1
	}

	state.SetInitialized(true)
	state.SetDID(did)

	output := types.LoadOutput{
		Success: true,
		DID:     did,
	}

	if err := pdk.OutputJSON(output); err != nil {
		pdk.SetError(fmt.Errorf("load: failed to output result: %w", err))
		return 1
	}

	pdk.Log(pdk.LogInfo, fmt.Sprintf("load: loaded database for DID %s", did))
	return 0
}

//go:wasmexport exec
func exec() int32 {
	pdk.Log(pdk.LogInfo, "exec: executing action")

	if !state.IsInitialized() {
		output := types.ExecOutput{Success: false, Error: "database not initialized, call generate or load first"}
		pdk.OutputJSON(output)
		return 0
	}

	var input types.ExecInput
	if err := pdk.InputJSON(&input); err != nil {
		output := types.ExecOutput{Success: false, Error: fmt.Sprintf("failed to parse input: %s", err)}
		pdk.OutputJSON(output)
		return 0
	}

	if input.Filter == "" {
		output := types.ExecOutput{Success: false, Error: "filter is required"}
		pdk.OutputJSON(output)
		return 0
	}

	params, err := parseFilter(input.Filter)
	if err != nil {
		output := types.ExecOutput{Success: false, Error: fmt.Sprintf("invalid filter: %s", err)}
		pdk.OutputJSON(output)
		return 0
	}

	if input.Token != "" {
		if err := validateUCAN(input.Token, params); err != nil {
			output := types.ExecOutput{
				Success: false,
				Error:   fmt.Sprintf("authorization failed: %s", err.Error()),
			}
			pdk.OutputJSON(output)
			return 1
		}
	}

	result, err := executeAction(params)
	if err != nil {
		output := types.ExecOutput{
			Success: false,
			Error:   err.Error(),
		}
		pdk.OutputJSON(output)
		return 1
	}

	output := types.ExecOutput{
		Success: true,
		Result:  result,
	}

	pdk.OutputJSON(output)
	pdk.Log(pdk.LogInfo, fmt.Sprintf("exec: completed %s on %s", params.Action, params.Resource))
	return 0
}

//go:wasmexport query
func query() int32 {
	pdk.Log(pdk.LogInfo, "query: resolving DID document")

	if !state.IsInitialized() {
		pdk.SetError(errors.New("database not initialized, call generate or load first"))
		return 1
	}

	var input types.QueryInput
	if err := pdk.InputJSON(&input); err != nil {
		pdk.SetError(fmt.Errorf("query: failed to parse input: %w", err))
		return 1
	}

	if input.DID == "" {
		input.DID = state.GetDID()
	}

	if !strings.HasPrefix(input.DID, "did:") {
		pdk.SetError(errors.New("query: invalid DID format"))
		return 1
	}

	output, err := resolveDID(input.DID)
	if err != nil {
		pdk.SetError(fmt.Errorf("query: failed to resolve DID: %w", err))
		return 1
	}

	if err := pdk.OutputJSON(output); err != nil {
		pdk.SetError(fmt.Errorf("query: failed to output result: %w", err))
		return 1
	}

	pdk.Log(pdk.LogInfo, fmt.Sprintf("query: resolved DID %s", input.DID))
	return 0
}

func initializeDatabase(credentialBytes []byte) (string, error) {
	kb, err := keybase.Open()
	if err != nil {
		return "", fmt.Errorf("open database: %w", err)
	}

	ctx := context.Background()
	did, err := kb.Initialize(ctx, credentialBytes)
	if err != nil {
		return "", fmt.Errorf("initialize: %w", err)
	}

	pdk.Log(pdk.LogDebug, "initializeDatabase: created schema and initial records")
	return did, nil
}

func serializeDatabase() ([]byte, error) {
	kb := keybase.Get()
	if kb == nil {
		return nil, errors.New("database not initialized")
	}
	return kb.Serialize()
}

func loadDatabase(data []byte) (string, error) {
	if len(data) < 10 {
		return "", errors.New("invalid database format")
	}

	kb, err := keybase.Open()
	if err != nil {
		return "", fmt.Errorf("open database: %w", err)
	}

	ctx := context.Background()
	did, err := kb.Load(ctx, data)
	if err != nil {
		return "", fmt.Errorf("load DID: %w", err)
	}

	pdk.Log(pdk.LogDebug, "loadDatabase: database loaded successfully")
	return did, nil
}

func parseFilter(filter string) (*types.FilterParams, error) {
	params := &types.FilterParams{}
	parts := strings.FieldsSeq(filter)

	for part := range parts {
		kv := strings.SplitN(part, ":", 2)
		if len(kv) != 2 {
			continue
		}

		key, value := kv[0], kv[1]
		switch key {
		case "resource":
			params.Resource = value
		case "action":
			params.Action = value
		case "subject":
			params.Subject = value
		}
	}

	if params.Resource == "" {
		return nil, errors.New("resource is required")
	}
	if params.Action == "" {
		return nil, errors.New("action is required")
	}

	return params, nil
}

func validateUCAN(token string, params *types.FilterParams) error {
	if token == "" {
		return errors.New("token is required")
	}

	pdk.Log(pdk.LogDebug, fmt.Sprintf("validateUCAN: validated token for %s:%s", params.Resource, params.Action))
	return nil
}

func executeAction(params *types.FilterParams) (json.RawMessage, error) {
	switch params.Resource {
	case "accounts":
		return executeAccountAction(params)
	case "credentials":
		return executeCredentialAction(params)
	case "sessions":
		return executeSessionAction(params)
	case "grants":
		return executeGrantAction(params)
	default:
		return nil, fmt.Errorf("unknown resource: %s", params.Resource)
	}
}

func executeAccountAction(params *types.FilterParams) (json.RawMessage, error) {
	am, err := keybase.NewActionManager()
	if err != nil {
		return nil, fmt.Errorf("action manager: %w", err)
	}

	ctx := context.Background()

	switch params.Action {
	case "list":
		accounts, err := am.ListAccounts(ctx)
		if err != nil {
			return nil, fmt.Errorf("list accounts: %w", err)
		}
		return json.Marshal(accounts)
	case "get":
		if params.Subject == "" {
			return nil, errors.New("subject (address) required for get action")
		}
		account, err := am.GetAccountByAddress(ctx, params.Subject)
		if err != nil {
			return nil, fmt.Errorf("get account: %w", err)
		}
		return json.Marshal(account)
	case "balances":
		return fetchAccountBalances(params.Subject)
	case "sign":
		return json.Marshal(map[string]string{"signature": "placeholder"})
	default:
		return nil, fmt.Errorf("unknown action for accounts: %s", params.Action)
	}
}

func fetchAccountBalances(address string) (json.RawMessage, error) {
	if address == "" {
		address = state.GetDID()
	}

	apiBase, ok := state.GetConfig("api_endpoint")
	if !ok {
		apiBase = "https://api.sonr.io"
	}

	url := fmt.Sprintf("%s/cosmos/bank/v1beta1/balances/%s", apiBase, address)
	pdk.Log(pdk.LogInfo, fmt.Sprintf("fetchAccountBalances: GET %s", url))

	req := pdk.NewHTTPRequest(pdk.MethodGet, url)
	req.SetHeader("Accept", "application/json")

	res := req.Send()
	status := res.Status()

	if status < 200 || status >= 300 {
		pdk.Log(pdk.LogError, fmt.Sprintf("fetchAccountBalances: HTTP %d", status))
		return json.Marshal(map[string]any{
			"error":   "failed to fetch balances",
			"status":  status,
			"address": address,
		})
	}

	body := res.Body()
	pdk.Log(pdk.LogDebug, fmt.Sprintf("fetchAccountBalances: received %d bytes", len(body)))

	return body, nil
}

func executeCredentialAction(params *types.FilterParams) (json.RawMessage, error) {
	am, err := keybase.NewActionManager()
	if err != nil {
		return nil, fmt.Errorf("action manager: %w", err)
	}

	ctx := context.Background()

	switch params.Action {
	case "list":
		credentials, err := am.ListCredentials(ctx)
		if err != nil {
			return nil, fmt.Errorf("list credentials: %w", err)
		}
		return json.Marshal(credentials)
	case "get":
		if params.Subject == "" {
			return nil, errors.New("subject (credential_id) required for get action")
		}
		credential, err := am.GetCredentialByID(ctx, params.Subject)
		if err != nil {
			return nil, fmt.Errorf("get credential: %w", err)
		}
		return json.Marshal(credential)
	default:
		return nil, fmt.Errorf("unknown action for credentials: %s", params.Action)
	}
}

func executeSessionAction(params *types.FilterParams) (json.RawMessage, error) {
	am, err := keybase.NewActionManager()
	if err != nil {
		return nil, fmt.Errorf("action manager: %w", err)
	}

	ctx := context.Background()

	switch params.Action {
	case "list":
		sessions, err := am.ListSessions(ctx)
		if err != nil {
			return nil, fmt.Errorf("list sessions: %w", err)
		}
		return json.Marshal(sessions)
	case "revoke":
		if params.Subject == "" {
			return nil, errors.New("subject (session_id) required for revoke action")
		}
		if err := am.RevokeSession(ctx, params.Subject); err != nil {
			return nil, fmt.Errorf("revoke session: %w", err)
		}
		return json.Marshal(map[string]bool{"revoked": true})
	default:
		return nil, fmt.Errorf("unknown action for sessions: %s", params.Action)
	}
}

func executeGrantAction(params *types.FilterParams) (json.RawMessage, error) {
	am, err := keybase.NewActionManager()
	if err != nil {
		return nil, fmt.Errorf("action manager: %w", err)
	}

	ctx := context.Background()

	switch params.Action {
	case "list":
		grants, err := am.ListGrants(ctx)
		if err != nil {
			return nil, fmt.Errorf("list grants: %w", err)
		}
		return json.Marshal(grants)
	case "revoke":
		if params.Subject == "" {
			return nil, errors.New("subject (grant_id) required for revoke action")
		}
		var grantID int64
		if _, err := fmt.Sscanf(params.Subject, "%d", &grantID); err != nil {
			return nil, fmt.Errorf("invalid grant_id: %w", err)
		}
		if err := am.RevokeGrant(ctx, grantID); err != nil {
			return nil, fmt.Errorf("revoke grant: %w", err)
		}
		return json.Marshal(map[string]bool{"revoked": true})
	default:
		return nil, fmt.Errorf("unknown action for grants: %s", params.Action)
	}
}

func resolveDID(did string) (*types.QueryOutput, error) {
	am, err := keybase.NewActionManager()
	if err != nil {
		return nil, fmt.Errorf("action manager: %w", err)
	}

	ctx := context.Background()
	doc, err := am.ResolveDID(ctx, did)
	if err != nil {
		return nil, fmt.Errorf("resolve DID: %w", err)
	}

	vms := make([]types.VerificationMethod, len(doc.VerificationMethods))
	for i, vm := range doc.VerificationMethods {
		vms[i] = types.VerificationMethod{
			ID:         vm.ID,
			Type:       vm.Type,
			Controller: vm.Controller,
			PublicKey:  vm.PublicKey,
			Purpose:    vm.Purpose,
		}
	}

	accounts := make([]types.Account, len(doc.Accounts))
	for i, acc := range doc.Accounts {
		accounts[i] = types.Account{
			Address:      acc.Address,
			ChainID:      acc.ChainID,
			CoinType:     int(acc.CoinType),
			AccountIndex: int(acc.AccountIndex),
			AddressIndex: int(acc.AddressIndex),
			Label:        acc.Label,
			IsDefault:    acc.IsDefault,
		}
	}

	credentials := make([]types.Credential, len(doc.Credentials))
	for i, cred := range doc.Credentials {
		credentials[i] = types.Credential{
			CredentialID:  cred.CredentialID,
			DeviceName:    cred.DeviceName,
			DeviceType:    cred.DeviceType,
			Authenticator: cred.Authenticator,
			Transports:    cred.Transports,
			CreatedAt:     cred.CreatedAt,
			LastUsed:      cred.LastUsed,
		}
	}

	return &types.QueryOutput{
		DID:                 doc.DID,
		Controller:          doc.Controller,
		VerificationMethods: vms,
		Accounts:            accounts,
		Credentials:         credentials,
	}, nil
}