CanvasArchiver/internal/auth/auth.go

package auth

import (
	"bufio"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strings"

	"git.directme.in/Joren/CanvasArchiver/internal/config"
	"git.directme.in/Joren/CanvasArchiver/internal/models"
)

type Authenticator struct {
	HTTPClient *http.Client
}

func NewAuthenticator(client *http.Client) *Authenticator {
	return &Authenticator{
		HTTPClient: client,
	}
}

func (a *Authenticator) GetToken() (string, error) {
	creds, err := LoadCredentials()
	if err == nil && creds.AuthMethod != "" {
		switch creds.AuthMethod {
		case "api":
			token := strings.TrimSpace(creds.APIToken)
			if token == "" {
				return a.promptMethod()
			}
			fmt.Print("[*] Validating saved API token... ")
			if err := a.validateAPIToken(token); err != nil {
				fmt.Println("FAILED")
				fmt.Printf("[!] %v\n", err)
				SaveCredentials(&models.Credentials{})
				return a.promptMethod()
			}
			fmt.Println("OK")
			return token, nil
		case "oauth":
			return a.refreshOrLogin()
		}
	}
	return a.promptMethod()
}

func (a *Authenticator) refreshOrLogin() (string, error) {
	creds, err := LoadCredentials()
	if err != nil {
		return a.login()
	}
	if strings.TrimSpace(creds.RefreshToken) == "" {
		return a.login()
	}

	fmt.Println("[*] Reusing saved refresh token...")
	tr, err := a.doTokenRequest(url.Values{
		"grant_type":    {"refresh_token"},
		"client_id":     {config.ClientID},
		"client_secret": {config.ClientSecret},
		"refresh_token": {creds.RefreshToken},
		"redirect_uri":  {config.RedirectURI},
	})
	if err != nil {
		fmt.Printf("[!] Refresh failed: %v\n", err)
		return a.login()
	}
	fmt.Println("[+] Session refreshed.")
	return tr.AccessToken, nil
}

func (a *Authenticator) promptMethod() (string, error) {
	fmt.Println("Select authentication method:")
	fmt.Println("  [1] Login via browser (OAuth)")
	fmt.Println("  [2] Use Canvas API token")
	fmt.Print("Choice (1/2): ")
	input, err := bufio.NewReader(os.Stdin).ReadString('\n')
	if err != nil && err != io.EOF {
		return "", err
	}
	switch strings.TrimSpace(input) {
	case "2":
		return a.setupAPIToken()
	default:
		return a.login()
	}
}

func (a *Authenticator) setupAPIToken() (string, error) {
	fmt.Print("Enter Canvas API token: ")
	token, err := bufio.NewReader(os.Stdin).ReadString('\n')
	if err != nil && err != io.EOF {
		return "", err
	}
	token = strings.TrimSpace(token)
	if token == "" {
		return "", fmt.Errorf("empty API token")
	}

	fmt.Print("[*] Validating API token... ")
	if err := a.validateAPIToken(token); err != nil {
		fmt.Println("FAILED")
		return "", err
	}
	fmt.Println("OK")

	if err := SaveCredentials(&models.Credentials{
		AuthMethod: "api",
		APIToken:   token,
	}); err != nil {
		return "", err
	}
	fmt.Println("[+] API token saved.")
	return token, nil
}

func (a *Authenticator) validateAPIToken(token string) error {
	req, _ := http.NewRequest("GET", config.BaseURL+"/api/v1/users/self", nil)
	req.Header.Set("Authorization", "Bearer "+token)
	resp, err := a.HTTPClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("invalid token: %d %s", resp.StatusCode, strings.TrimSpace(string(body)))
	}
	return nil
}

func (a *Authenticator) login() (string, error) {
	codeVerifier, codeChallenge, err := generatePKCEPair()
	if err != nil {
		return "", err
	}

	fmt.Println("--- Initial Canvas Login Required ---")
	fmt.Printf("Visit: %s\n", buildAuthURL(codeChallenge))
	fmt.Print("Enter Code or redirect URL: ")
	input, err := bufio.NewReader(os.Stdin).ReadString('\n')
	if err != nil && err != io.EOF {
		return "", err
	}

	code, err := extractCode(input)
	if err != nil {
		return "", err
	}

	tr, err := a.doTokenRequest(url.Values{
		"grant_type":    {"authorization_code"},
		"client_id":     {config.ClientID},
		"client_secret": {config.ClientSecret},
		"redirect_uri":  {config.RedirectURI},
		"code":          {code},
		"code_verifier": {codeVerifier},
	})
	if err != nil {
		return "", err
	}

	if err := SaveRefreshToken(tr.RefreshToken); err != nil {
		return "", err
	}
	fmt.Println("[+] Login successful.")
	return tr.AccessToken, nil
}

func (a *Authenticator) doTokenRequest(v url.Values) (*models.TokenResponse, error) {
	resp, err := a.HTTPClient.PostForm(config.BaseURL+"/login/oauth2/token", v)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("token request failed: %d: %s", resp.StatusCode, string(body))
	}

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

func buildAuthURL(codeChallenge string) string {
	u, _ := url.Parse(config.BaseURL + "/login/oauth2/auth")
	q := u.Query()
	q.Set("client_id", config.ClientID)
	q.Set("response_type", "code")
	q.Set("mobile", "1")
	q.Set("purpose", config.OAuthPurpose)
	q.Set("code_challenge", codeChallenge)
	q.Set("code_challenge_method", "S256")
	q.Set("redirect_uri", config.RedirectURI)
	if config.AuthProvider != "" {
		q.Set("authentication_provider", config.AuthProvider)
	}
	u.RawQuery = q.Encode()
	return u.String()
}

func generatePKCEPair() (string, string, error) {
	randomBytes := make([]byte, 64)
	if _, err := rand.Read(randomBytes); err != nil {
		return "", "", err
	}

	verifier := base64.RawURLEncoding.EncodeToString(randomBytes)
	challengeBytes := sha256.Sum256([]byte(verifier))
	challenge := base64.RawURLEncoding.EncodeToString(challengeBytes[:])
	return verifier, challenge, nil
}

func extractCode(input string) (string, error) {
	input = strings.TrimSpace(input)
	if input == "" {
		return "", fmt.Errorf("empty authorization code")
	}

	for _, field := range strings.Fields(input) {
		if strings.Contains(field, "code=") {
			input = field
			break
		}
	}

	parsedURL, err := url.Parse(input)
	if err == nil {
		if code := parsedURL.Query().Get("code"); code != "" {
			return code, nil
		}
	}

	if values, err := url.ParseQuery(input); err == nil {
		if code := values.Get("code"); code != "" {
			return code, nil
		}
	}

	return input, nil
}

func SaveCredentials(creds *models.Credentials) error {
	data, err := json.MarshalIndent(creds, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(config.CredsFile, data, 0o644)
}

func SaveRefreshToken(refreshToken string) error {
	creds, err := LoadCredentials()
	if err != nil {
		creds = &models.Credentials{}
	}
	creds.RefreshToken = refreshToken
	creds.AuthMethod = "oauth"
	return SaveCredentials(creds)
}

func LoadCredentials() (*models.Credentials, error) {
	data, err := os.ReadFile(config.CredsFile)
	if err != nil {
		return nil, err
	}
	var creds models.Credentials
	if err := json.Unmarshal(data, &creds); err != nil {
		return nil, err
	}
	return &creds, nil
}