go-eve-pi/esi_sso.go

package main

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
)

const (
	issuerAuthorizeURL = "https://login.eveonline.com/v2/oauth/authorize"
	issuerTokenURL     = "https://login.eveonline.com/v2/oauth/token"
)

type Token struct {
	ID            uint   `gorm:"primaryKey"`
	CharacterName string `gorm:"uniqueIndex"`
	AccessToken   string
	RefreshToken  string
	ExpiresAt     time.Time
	UpdatedAt     time.Time
	CreatedAt     time.Time
}

type SSO struct {
	clientID     string
	redirectURI  string
	scopes       []string
	db           *gorm.DB
	mu           sync.Mutex
	server       *http.Server
	state        string
	callbackChan chan struct {
		code  string
		state string
		err   error
	}
}

// NewSSO creates a new SSO instance
func NewSSO(clientID, redirectURI string, scopes []string) (*SSO, error) {
	s := &SSO{
		clientID:    clientID,
		redirectURI: redirectURI,
		scopes:      scopes,
	}

	if err := s.initDB(); err != nil {
		return nil, err
	}

	return s, nil
}

func (s *SSO) initDB() error {
	home, err := os.UserHomeDir()
	if err != nil {
		return err
	}
	dbPath := filepath.Join(home, ".industrializer", "sqlite-latest.sqlite")
	db, err := gorm.Open(sqlite.Open(dbPath), &gorm.Config{})
	if err != nil {
		return err
	}
	if err := db.AutoMigrate(&Token{}); err != nil {
		return err
	}
	s.db = db
	return nil
}

// GetToken returns a valid access token for the given character name
// If no token exists, it will start the OAuth flow
// If token is expired, it will refresh it automatically
func (s *SSO) GetToken(ctx context.Context, characterName string) (string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Try to get existing token from DB
	var token Token
	if err := s.db.Where("character_name = ?", characterName).First(&token).Error; err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			// No token exists, need to authenticate
			if err := s.startAuthFlow(ctx, characterName); err != nil {
				return "", err
			}
			// After authentication, fetch the token from DB
			if err := s.db.Where("character_name = ?", characterName).First(&token).Error; err != nil {
				return "", err
			}
		} else {
			return "", err
		}
	}

	// Check if token needs refresh
	if time.Now().After(token.ExpiresAt.Add(-60 * time.Second)) {
		if err := s.refreshToken(ctx, &token); err != nil {
			// Refresh failed, need to re-authenticate
			if err := s.startAuthFlow(ctx, characterName); err != nil {
				return "", err
			}
			// After re-authentication, fetch the token from DB
			if err := s.db.Where("character_name = ?", characterName).First(&token).Error; err != nil {
				return "", err
			}
		}
	}

	return token.AccessToken, nil
}

func (s *SSO) startAuthFlow(ctx context.Context, characterName string) error {
	// Generate PKCE
	verifier, challenge, err := generatePKCE()
	if err != nil {
		return err
	}

	s.state = randString(24)
	s.callbackChan = make(chan struct {
		code  string
		state string
		err   error
	}, 1)
	authURL := s.buildAuthURL(challenge, s.state)

	fmt.Printf("Please visit this URL to authenticate:\n%s\n", authURL)
	fmt.Println("Waiting for authentication...")

	// Start callback server
	server, err := s.startCallbackServer()
	if err != nil {
		return err
	}
	s.server = server
	defer server.Shutdown(ctx)

	// Wait for callback
	code, receivedState, err := s.waitForCallback()
	if err != nil {
		return err
	}

	if receivedState != s.state {
		return errors.New("invalid state parameter")
	}

	// Exchange code for token
	token, err := s.exchangeCodeForToken(ctx, code, verifier)
	if err != nil {
		return err
	}

	// Save token to DB
	token.CharacterName = characterName
	return s.db.Save(&token).Error
}

func (s *SSO) buildAuthURL(challenge, state string) string {
	q := url.Values{}
	q.Set("response_type", "code")
	q.Set("client_id", s.clientID)
	q.Set("redirect_uri", s.redirectURI)
	if len(s.scopes) > 0 {
		q.Set("scope", strings.Join(s.scopes, " "))
	}
	q.Set("state", state)
	q.Set("code_challenge", challenge)
	q.Set("code_challenge_method", "S256")

	return issuerAuthorizeURL + "?" + q.Encode()
}

func (s *SSO) startCallbackServer() (*http.Server, error) {
	u, err := url.Parse(s.redirectURI)
	if err != nil {
		return nil, err
	}
	if u.Scheme != "http" && u.Scheme != "https" {
		return nil, errors.New("redirect URI must be http(s)")
	}
	hostPort := u.Host
	if !strings.Contains(hostPort, ":") {
		if u.Scheme == "https" {
			hostPort += ":443"
		} else {
			hostPort += ":80"
		}
	}

	mux := http.NewServeMux()
	mux.HandleFunc(u.Path, func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodGet {
			w.WriteHeader(http.StatusMethodNotAllowed)
			s.callbackChan <- struct {
				code  string
				state string
				err   error
			}{"", "", errors.New("method not allowed")}
			return
		}
		q := r.URL.Query()
		code := q.Get("code")
		st := q.Get("state")
		if code == "" || st == "" || st != s.state {
			fmt.Printf("Invalid SSO response: code=%s, state=%s, expected_state=%s\n", code, st, s.state)
			w.WriteHeader(http.StatusBadRequest)
			_, _ = w.Write([]byte("Invalid SSO response"))
			s.callbackChan <- struct {
				code  string
				state string
				err   error
			}{"", "", errors.New("invalid state")}
			return
		}
		fmt.Printf("Exchanging token for code: %s\n", code)
		w.Header().Set("Content-Type", "text/html")
		_, _ = w.Write([]byte("<html><body><h1>Login successful!</h1><p>You can close this window.</p></body></html>"))
		s.callbackChan <- struct {
			code  string
			state string
			err   error
		}{code, st, nil}
		go func() {
			time.Sleep(200 * time.Millisecond)
			_ = s.server.Shutdown(context.Background())
		}()
	})

	ln, err := net.Listen("tcp", hostPort)
	if err != nil {
		return nil, err
	}

	server := &http.Server{Handler: mux}
	go func() { _ = server.Serve(ln) }()
	return server, nil
}

func (s *SSO) waitForCallback() (code, state string, err error) {
	// Wait for callback through channel
	select {
	case result := <-s.callbackChan:
		return result.code, result.state, result.err
	case <-time.After(30 * time.Second):
		return "", "", errors.New("callback timeout")
	}
}

func (s *SSO) exchangeCodeForToken(ctx context.Context, code, verifier string) (*Token, error) {
	form := url.Values{}
	form.Set("grant_type", "authorization_code")
	form.Set("code", code)
	form.Set("client_id", s.clientID)
	form.Set("code_verifier", verifier)

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, issuerTokenURL, strings.NewReader(form.Encode()))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		b, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("token exchange failed: %s: %s", resp.Status, string(b))
	}

	var tr tokenResponse
	if err := json.NewDecoder(resp.Body).Decode(&tr); err != nil {
		return nil, err
	}

	// Parse character info from token
	name, _ := parseTokenCharacter(tr.AccessToken)

	return &Token{
		CharacterName: name,
		AccessToken:   tr.AccessToken,
		RefreshToken:  tr.RefreshToken,
		ExpiresAt:     time.Now().Add(time.Duration(tr.ExpiresIn-30) * time.Second),
	}, nil
}

func (s *SSO) refreshToken(ctx context.Context, token *Token) error {
	form := url.Values{}
	form.Set("grant_type", "refresh_token")
	form.Set("refresh_token", token.RefreshToken)
	form.Set("client_id", s.clientID)

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, issuerTokenURL, strings.NewReader(form.Encode()))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		b, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("token refresh failed: %s: %s", resp.Status, string(b))
	}

	var tr tokenResponse
	if err := json.NewDecoder(resp.Body).Decode(&tr); err != nil {
		return err
	}

	// Update token
	token.AccessToken = tr.AccessToken
	if tr.RefreshToken != "" {
		token.RefreshToken = tr.RefreshToken
	}
	if tr.ExpiresIn > 0 {
		token.ExpiresAt = time.Now().Add(time.Duration(tr.ExpiresIn-30) * time.Second)
	}

	return s.db.Save(token).Error
}

// Utility functions
func generatePKCE() (verifier string, challenge string, err error) {
	buf := make([]byte, 32)
	if _, err = rand.Read(buf); err != nil {
		return
	}
	v := base64.RawURLEncoding.EncodeToString(buf)
	h := sha256.Sum256([]byte(v))
	c := base64.RawURLEncoding.EncodeToString(h[:])
	return v, c, nil
}

func randString(n int) string {
	buf := make([]byte, n)
	_, _ = rand.Read(buf)
	return base64.RawURLEncoding.EncodeToString(buf)
}

func parseTokenCharacter(jwt string) (name string, id int64) {
	parts := strings.Split(jwt, ".")
	if len(parts) != 3 {
		return "", 0
	}
	payload, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		return "", 0
	}
	var m map[string]any
	if err := json.Unmarshal(payload, &m); err != nil {
		return "", 0
	}
	if v, ok := m["name"].(string); ok {
		name = v
	}
	if v, ok := m["sub"].(string); ok {
		if idx := strings.LastIndexByte(v, ':'); idx > -1 {
			if idv, err := strconv.ParseInt(v[idx+1:], 10, 64); err == nil {
				id = idv
			}
		}
	}
	return
}

type tokenResponse struct {
	AccessToken  string `json:"access_token"`
	RefreshToken string `json:"refresh_token"`
	TokenType    string `json:"token_type"`
	ExpiresIn    int    `json:"expires_in"`
}