torrent-gateway/internal/auth/nostr_auth.go

package auth

import (
	"context"
	"crypto/rand"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"log"
	mathrand "math/rand"
	"net/url"
	"strings"
	"time"

	"github.com/nbd-wtf/go-nostr"
	"github.com/nbd-wtf/go-nostr/nip19"
	"github.com/nbd-wtf/go-nostr/nip44"
)

// NostrAuth handles Nostr-based authentication
type NostrAuth struct {
	db *sql.DB
}

// NewNostrAuth creates a new Nostr authentication handler
func NewNostrAuth(db *sql.DB) *NostrAuth {
	return &NostrAuth{db: db}
}

// User represents a user in the system
type User struct {
	Pubkey       string    `json:"pubkey"`
	DisplayName  string    `json:"display_name"`
	ProfileImage string    `json:"profile_image"`
	CreatedAt    time.Time `json:"created_at"`
	LastLogin    time.Time `json:"last_login"`
	StorageUsed  int64     `json:"storage_used"`
	FileCount    int       `json:"file_count"`
}

// Session represents an active user session
type Session struct {
	Token     string    `json:"token"`
	Pubkey    string    `json:"pubkey"`
	CreatedAt time.Time `json:"created_at"`
	ExpiresAt time.Time `json:"expires_at"`
}

// AuthEvent represents a Nostr authentication event
type AuthEvent struct {
	Event     *nostr.Event `json:"event"`
	Challenge string       `json:"challenge,omitempty"`
}

// ValidateNIP07 validates a NIP-07 authentication event
func (na *NostrAuth) ValidateNIP07(authEventJSON string) (string, error) {
	var authEvent AuthEvent
	if err := json.Unmarshal([]byte(authEventJSON), &authEvent); err != nil {
		return "", fmt.Errorf("invalid auth event JSON: %w", err)
	}

	if authEvent.Event == nil {
		return "", fmt.Errorf("missing event in auth data")
	}

	event := authEvent.Event

	// For NIP-07, we can accept any kind of signed event as proof of key ownership
	// The standard approach is to use kind 22242 for auth events, but many implementations vary
	if event.Kind != 22242 && event.Kind != 27235 {
		log.Printf("Warning: Non-standard auth event kind %d, accepting anyway", event.Kind)
	}

	// Validate event timestamp (should be recent)
	now := time.Now()
	eventTime := time.Unix(int64(event.CreatedAt), 0)
	if now.Sub(eventTime) > 10*time.Minute { // More lenient time window
		return "", fmt.Errorf("event too old: %v", eventTime)
	}
	if eventTime.After(now.Add(2 * time.Minute)) {
		return "", fmt.Errorf("event from future: %v", eventTime)
	}

	// Validate signature
	if ok, err := event.CheckSignature(); !ok || err != nil {
		return "", fmt.Errorf("invalid signature: %v", err)
	}

	// Extract and validate challenge from tags if present
	var challenge string
	for _, tag := range event.Tags {
		if len(tag) >= 2 && tag[0] == "challenge" {
			challenge = tag[1]
			break
		}
	}

	// If challenge was provided in the auth event, validate it matches
	if authEvent.Challenge != "" && challenge != authEvent.Challenge {
		return "", fmt.Errorf("challenge mismatch")
	}

	return event.PubKey, nil
}

// ValidateNIP46 validates a NIP-46 bunker URL and returns pubkey
func (na *NostrAuth) ValidateNIP46(bunkerURL string) (string, error) {
	// Parse bunker URL format: bunker://<pubkey>?relay=<relay>&secret=<secret>
	// or nostrconnect://<pubkey>?relay=<relay>&metadata=<metadata>
	if !strings.HasPrefix(bunkerURL, "bunker://") && !strings.HasPrefix(bunkerURL, "nostrconnect://") {
		return "", fmt.Errorf("invalid bunker URL format, expected bunker:// or nostrconnect://")
	}

	parsedURL, err := url.Parse(bunkerURL)
	if err != nil {
		return "", fmt.Errorf("failed to parse bunker URL: %w", err)
	}

	pubkey := parsedURL.Host
	if pubkey == "" {
		return "", fmt.Errorf("missing pubkey in bunker URL")
	}

	// Validate pubkey format (should be hex)
	if len(pubkey) != 64 {
		return "", fmt.Errorf("invalid pubkey length: expected 64 chars, got %d", len(pubkey))
	}

	for _, c := range pubkey {
		if !((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F')) {
			return "", fmt.Errorf("invalid pubkey format: must be hex")
		}
	}

	// Extract relays and secret
	params := parsedURL.Query()
	relays := params["relay"]
	if len(relays) == 0 {
		return "", fmt.Errorf("no relays specified in bunker URL")
	}

	secret := ""
	if secrets := params["secret"]; len(secrets) > 0 {
		secret = secrets[0]
	}

	// Establish full NIP-46 connection
	return na.establishNIP46Connection(pubkey, relays, secret)
}

// establishNIP46Connection performs the full NIP-46 handshake
func (na *NostrAuth) establishNIP46Connection(remotePubkey string, relays []string, secret string) (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()

	// Generate client keypair for this connection
	clientSK := nostr.GeneratePrivateKey()
	clientPK, _ := nostr.GetPublicKey(clientSK)

	log.Printf("Starting NIP-46 connection to %s via relays %v", remotePubkey, relays)

	// Create relay pool
	pool := nostr.NewSimplePool(ctx)
	
	// Give relays time to connect
	time.Sleep(2 * time.Second)

	// Subscribe to responses from remote signer
	since := nostr.Timestamp(time.Now().Add(-1 * time.Minute).Unix())
	filters := []nostr.Filter{{
		Kinds: []int{24133}, // NIP-46 response events
		Tags: nostr.TagMap{
			"p": []string{clientPK}, // Events directed to our client
		},
		Since: &since,
	}}

	responseChan := make(chan *nostr.Event, 10)
	sub := pool.SubMany(ctx, relays, filters)
	
	// Listen for events in a goroutine
	go func() {
		for evt := range sub {
			// Only process events from the remote signer
			if evt.Event.PubKey == remotePubkey {
				responseChan <- evt.Event
			}
		}
	}()

	// Step 1: Send connect request
	connectID := generateRandomString(16)
	var connectParams []string
	if secret != "" {
		connectParams = []string{remotePubkey, secret}
	} else {
		connectParams = []string{remotePubkey}
	}
	
	connectReq := map[string]interface{}{
		"id":     connectID,
		"method": "connect", 
		"params": connectParams,
	}

	if err := na.sendNIP46Request(ctx, pool, relays, clientSK, remotePubkey, connectReq); err != nil {
		return "", fmt.Errorf("failed to send connect request: %w", err)
	}

	log.Printf("Sent NIP-46 connect request: %+v", connectReq)
	log.Printf("Client pubkey: %s", clientPK)
	log.Printf("Remote pubkey: %s", remotePubkey)
	log.Printf("Waiting for response...")

	// Step 2: Wait for connect response, then send get_public_key
	var userPubkey string
	connectAcked := false
	getPkID := ""

	for {
		select {
		case <-ctx.Done():
			return "", fmt.Errorf("timeout waiting for remote signer response")
		case evt := <-responseChan:
			// Decrypt the response
			response, err := na.decryptNIP46Response(clientSK, remotePubkey, evt)
			if err != nil {
				log.Printf("Failed to decrypt NIP-46 response: %v", err)
				continue
			}

			log.Printf("Received NIP-46 response: %+v", response)

			// Handle connect response
			if responseID, ok := response["id"].(string); ok && responseID == connectID {
				if result, ok := response["result"].(string); ok && result == "ack" {
					connectAcked = true
					log.Printf("NIP-46 connect acknowledged")

					// Send get_public_key request
					getPkID = generateRandomString(16)
					getPkReq := map[string]interface{}{
						"id":     getPkID,
						"method": "get_public_key",
						"params": []string{},
					}

					if err := na.sendNIP46Request(ctx, pool, relays, clientSK, remotePubkey, getPkReq); err != nil {
						return "", fmt.Errorf("failed to send get_public_key request: %w", err)
					}

					log.Printf("Sent get_public_key request, waiting for user approval...")
				} else if errorMsg, ok := response["error"].(string); ok {
					return "", fmt.Errorf("connect request failed: %s", errorMsg)
				}
			}

			// Handle get_public_key response
			if connectAcked && getPkID != "" {
				if responseID, ok := response["id"].(string); ok && responseID == getPkID {
					if result, ok := response["result"].(string); ok {
						userPubkey = result
						log.Printf("Received user public key: %s", userPubkey)
						return userPubkey, nil
					} else if errorMsg, ok := response["error"].(string); ok {
						return "", fmt.Errorf("get_public_key request failed: %s", errorMsg)
					}
				}
			}
		}
	}
}

// sendNIP46Request sends an encrypted NIP-46 request
func (na *NostrAuth) sendNIP46Request(ctx context.Context, pool *nostr.SimplePool, relays []string, clientSK, remotePubkey string, request map[string]interface{}) error {
	// Serialize request
	requestJSON, err := json.Marshal(request)
	if err != nil {
		return fmt.Errorf("failed to marshal request: %w", err)
	}

	log.Printf("Sending NIP-46 request JSON: %s", string(requestJSON))

	// Encrypt request using NIP-44
	// Ensure remotePubkey is in correct hex format (no 02 prefix)
	cleanRemotePubkey := remotePubkey
	if len(remotePubkey) == 66 && strings.HasPrefix(remotePubkey, "02") {
		cleanRemotePubkey = remotePubkey[2:]
	}
	
	conversationKey, err := nip44.GenerateConversationKey(clientSK, cleanRemotePubkey)
	if err != nil {
		return fmt.Errorf("failed to generate conversation key: %w", err)
	}

	encryptedContent, err := nip44.Encrypt(string(requestJSON), conversationKey)
	if err != nil {
		return fmt.Errorf("failed to encrypt request: %w", err)
	}

	log.Printf("Encrypted content length: %d", len(encryptedContent))

	// Create event
	clientPK, _ := nostr.GetPublicKey(clientSK)
	evt := nostr.Event{
		Kind:      24133,
		CreatedAt: nostr.Now(),
		Tags: nostr.Tags{
			{"p", remotePubkey},
			{"relay", relays[0]}, // Add relay tag
		},
		Content: encryptedContent,
		PubKey:  clientPK,
	}

	log.Printf("Created NIP-46 event: kind=%d, from=%s, to=%s, content_len=%d", 
		evt.Kind, clientPK, remotePubkey, len(encryptedContent))

	// Sign event
	if err := evt.Sign(clientSK); err != nil {
		return fmt.Errorf("failed to sign event: %w", err)
	}

	// Publish to all relays
	for _, relayURL := range relays {
		relay, err := pool.EnsureRelay(relayURL)
		if err != nil {
			log.Printf("Failed to connect to relay %s: %v", relayURL, err)
			continue
		}

		log.Printf("Connected to relay %s, publishing event...", relayURL)
		
		if err := relay.Publish(ctx, evt); err != nil {
			log.Printf("Failed to publish to relay %s: %v", relayURL, err)
		} else {
			log.Printf("Published NIP-46 request to relay %s (event ID: %s)", relayURL, evt.ID)
		}
	}

	return nil
}

// decryptNIP46Response decrypts a NIP-46 response event
func (na *NostrAuth) decryptNIP46Response(clientSK, remotePubkey string, evt *nostr.Event) (map[string]interface{}, error) {
	// Ensure remotePubkey is in correct hex format (no 02 prefix)
	cleanRemotePubkey := remotePubkey
	if len(remotePubkey) == 66 && strings.HasPrefix(remotePubkey, "02") {
		cleanRemotePubkey = remotePubkey[2:]
	}
	
	// Generate conversation key
	conversationKey, err := nip44.GenerateConversationKey(clientSK, cleanRemotePubkey)
	if err != nil {
		return nil, fmt.Errorf("failed to generate conversation key: %w", err)
	}

	// Decrypt content
	decryptedJSON, err := nip44.Decrypt(evt.Content, conversationKey)
	if err != nil {
		return nil, fmt.Errorf("failed to decrypt content: %w", err)
	}

	// Parse JSON response
	var response map[string]interface{}
	if err := json.Unmarshal([]byte(decryptedJSON), &response); err != nil {
		return nil, fmt.Errorf("failed to parse response JSON: %w", err)
	}

	return response, nil
}

// generateRandomString generates a random string of specified length
func generateRandomString(length int) string {
	const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
	b := make([]byte, length)
	for i := range b {
		b[i] = charset[mathrand.Intn(len(charset))]
	}
	return string(b)
}

// CreateSession creates a new session for the given pubkey
func (na *NostrAuth) CreateSession(pubkey string) (string, error) {
	// Generate random session token
	tokenBytes := make([]byte, 32)
	if _, err := rand.Read(tokenBytes); err != nil {
		return "", fmt.Errorf("failed to generate session token: %w", err)
	}
	token := hex.EncodeToString(tokenBytes)

	// Session expires in 24 hours
	expiresAt := time.Now().Add(24 * time.Hour)

	// Store session in database
	_, err := na.db.Exec(`
		INSERT INTO sessions (token, pubkey, created_at, expires_at)
		VALUES (?, ?, ?, ?)
	`, token, pubkey, time.Now(), expiresAt)
	if err != nil {
		return "", fmt.Errorf("failed to store session: %w", err)
	}

	// Update user last login
	_, err = na.db.Exec(`
		INSERT INTO users (pubkey, last_login, created_at) 
		VALUES (?, ?, ?)
		ON CONFLICT(pubkey) DO UPDATE SET last_login = ?
	`, pubkey, time.Now(), time.Now(), time.Now())
	if err != nil {
		log.Printf("Warning: failed to update user login time: %v", err)
	}

	return token, nil
}

// ValidateSession validates a session token and returns the pubkey
func (na *NostrAuth) ValidateSession(token string) (string, error) {
	var session Session
	err := na.db.QueryRow(`
		SELECT token, pubkey, created_at, expires_at
		FROM sessions 
		WHERE token = ? AND expires_at > ?
	`, token, time.Now()).Scan(
		&session.Token, &session.Pubkey, 
		&session.CreatedAt, &session.ExpiresAt,
	)
	if err != nil {
		if err == sql.ErrNoRows {
			return "", fmt.Errorf("invalid or expired session")
		}
		return "", fmt.Errorf("failed to validate session: %w", err)
	}

	return session.Pubkey, nil
}

// GetUser retrieves user information by pubkey
func (na *NostrAuth) GetUser(pubkey string) (*User, error) {
	var user User
	err := na.db.QueryRow(`
		SELECT pubkey, COALESCE(display_name, ''), COALESCE(profile_image, ''),
		       created_at, last_login, COALESCE(storage_used, 0), COALESCE(file_count, 0)
		FROM users WHERE pubkey = ?
	`, pubkey).Scan(
		&user.Pubkey, &user.DisplayName, &user.ProfileImage,
		&user.CreatedAt, &user.LastLogin, &user.StorageUsed, &user.FileCount,
	)
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, nil
		}
		return nil, fmt.Errorf("failed to get user: %w", err)
	}

	return &user, nil
}

// UpdateUserProfile updates user profile information
func (na *NostrAuth) UpdateUserProfile(pubkey, displayName, profileImage string) error {
	_, err := na.db.Exec(`
		INSERT INTO users (pubkey, display_name, profile_image, created_at, last_login)
		VALUES (?, ?, ?, ?, ?)
		ON CONFLICT(pubkey) DO UPDATE SET 
		display_name = ?, profile_image = ?
	`, pubkey, displayName, profileImage, time.Now(), time.Now(), displayName, profileImage)
	if err != nil {
		return fmt.Errorf("failed to update user profile: %w", err)
	}

	return nil
}

// UpdateUserStats updates user storage statistics
func (na *NostrAuth) UpdateUserStats(pubkey string, storageUsed int64, fileCount int) error {
	_, err := na.db.Exec(`
		INSERT INTO users (pubkey, storage_used, file_count, created_at, last_login)
		VALUES (?, ?, ?, ?, ?)
		ON CONFLICT(pubkey) DO UPDATE SET 
		storage_used = ?, file_count = ?
	`, pubkey, storageUsed, fileCount, time.Now(), time.Now(), storageUsed, fileCount)
	if err != nil {
		return fmt.Errorf("failed to update user stats: %w", err)
	}

	return nil
}

// RevokeSession removes a session from the database
func (na *NostrAuth) RevokeSession(token string) error {
	_, err := na.db.Exec(`DELETE FROM sessions WHERE token = ?`, token)
	if err != nil {
		return fmt.Errorf("failed to revoke session: %w", err)
	}

	return nil
}

// CleanExpiredSessions removes expired sessions from the database
func (na *NostrAuth) CleanExpiredSessions() error {
	result, err := na.db.Exec(`DELETE FROM sessions WHERE expires_at < ?`, time.Now())
	if err != nil {
		return fmt.Errorf("failed to clean expired sessions: %w", err)
	}

	rowsAffected, _ := result.RowsAffected()
	if rowsAffected > 0 {
		log.Printf("Cleaned %d expired sessions", rowsAffected)
	}

	return nil
}

// GenerateChallenge generates a random challenge for authentication
func GenerateChallenge() (string, error) {
	challengeBytes := make([]byte, 16)
	if _, err := rand.Read(challengeBytes); err != nil {
		return "", fmt.Errorf("failed to generate challenge: %w", err)
	}
	return hex.EncodeToString(challengeBytes), nil
}

// ParsePubkeyFromNpub converts npub format to hex pubkey
func ParsePubkeyFromNpub(npub string) (string, error) {
	if !strings.HasPrefix(npub, "npub1") {
		return npub, nil // Already hex format
	}

	_, pubkeyBytes, err := nip19.Decode(npub)
	if err != nil {
		return "", fmt.Errorf("failed to decode npub: %w", err)
	}

	return hex.EncodeToString(pubkeyBytes.([]byte)), nil
}

// FormatPubkeyAsNpub converts hex pubkey to npub format
func FormatPubkeyAsNpub(pubkey string) (string, error) {
	pubkeyBytes, err := hex.DecodeString(pubkey)
	if err != nil {
		return "", fmt.Errorf("failed to decode pubkey: %w", err)
	}

	npub, err := nip19.EncodePublicKey(string(pubkeyBytes))
	if err != nil {
		return "", fmt.Errorf("failed to encode npub: %w", err)
	}

	return npub, nil
}