feat: add Okta MFA challenge flow for okta_saml

README.md

@@ -211,7 +211,7 @@ contexts:
 | `console_login` | Run `aws login` during `unic context setup`, then use the resulting profile-backed console credentials | `profile` |
 | `assume_role` | Assume a role from a base profile | `profile`, `role_arn` |
 | `sso` | Use AWS IAM Identity Center / SSO, reusing a valid AWS CLI SSO cache and prompting for login only when needed | `sso_start_url`, and for concrete contexts `sso_account_id`, `sso_role_name`; `profile` is optional |
-| `okta_saml` | Okta SAML federation: `unic env` signs in to Okta (prompt on stderr, password without echo; `UNIC_OKTA_USERNAME`/`UNIC_OKTA_PASSWORD` for automation), exchanges the SAML assertion via `sts:AssumeRoleWithSAML`, and caches the session under `~/.config/unic/cache/okta-saml/`. The TUI reuses a valid cached session passively. Okta MFA challenges are not supported yet | `okta_org_url`, `okta_app_id`; `role_arn` required when the assertion carries multiple roles. Passwords and MFA secrets are never stored |
+| `okta_saml` | Okta SAML federation: `unic env` signs in to Okta (prompt on stderr, password without echo; `UNIC_OKTA_USERNAME`/`UNIC_OKTA_PASSWORD` for automation), exchanges the SAML assertion via `sts:AssumeRoleWithSAML`, and caches the session under `~/.config/unic/cache/okta-saml/`. The TUI reuses a valid cached session passively. Okta MFA challenges support TOTP codes and Okta Verify push in v1 | `okta_org_url`, `okta_app_id`; `role_arn` required when the assertion carries multiple roles. Passwords and MFA secrets are never stored |
 
 The preferred context format separates `auth` from `resources`. `auth.sso_region` controls IAM Identity Center login and role-credential retrieval. `resources.default_region` is selected at startup, and `resources.regions` lists additional regions available from the global `R` region picker. Switching regions reuses the current credentials and recreates only the regional AWS clients.
 

internal/auth/okta.go

@@ -136,7 +136,23 @@ func promptOktaCredentials(orgURL string) (string, string, error) {
 
 type oktaAuthnResponse struct {
 	Status       string `json:"status"`
+	StateToken   string `json:"stateToken"`
 	SessionToken string `json:"sessionToken"`
+	FactorResult string `json:"factorResult"`
+	Embedded     struct {
+		Factors []oktaFactor `json:"factors"`
+	} `json:"_embedded"`
+}
+
+type oktaFactor struct {
+	ID         string `json:"id"`
+	FactorType string `json:"factorType"`
+	Provider   string `json:"provider"`
+	Links      struct {
+		Verify struct {
+			Href string `json:"href"`
+		} `json:"verify"`
+	} `json:"_links"`
 }
 
 // oktaPrimaryAuth performs Okta primary authentication and returns a one-time
@@ -177,13 +193,162 @@ func oktaPrimaryAuth(ctx context.Context, client *http.Client, orgURL, username,
 			return "", fmt.Errorf("okta returned SUCCESS without a session token")
 		}
 		return authn.SessionToken, nil
-	case "MFA_REQUIRED", "MFA_ENROLL":
-		return "", fmt.Errorf("okta requires an MFA challenge, which unic does not support yet (see issue #87)")
+	case "MFA_REQUIRED":
+		return oktaMFAChallenge(ctx, client, authn)
+	case "MFA_ENROLL":
+		return "", fmt.Errorf("okta account has no enrolled MFA factor; enroll one in Okta first")
 	default:
 		return "", fmt.Errorf("okta authentication ended in status %q", authn.Status)
 	}
 }
 
+// v1 MFA factor set: TOTP (token:software:totp) and Okta Verify push. TOTP is
+// preferred because it completes without waiting; other factor types are
+// rejected with an explicit list.
+const (
+	oktaFactorTOTP = "token:software:totp"
+	oktaFactorPush = "push"
+)
+
+var (
+	promptOktaMFACodeFn  = promptOktaMFACode
+	oktaPushPollInterval = 3 * time.Second
+	oktaPushPollTimeout  = 60 * time.Second
+)
+
+func promptOktaMFACode(factor oktaFactor) (string, error) {
+	fmt.Fprintf(os.Stderr, "Okta MFA code (%s %s): ", factor.Provider, factor.FactorType)
+	var code string
+	if _, err := fmt.Fscanln(os.Stdin, &code); err != nil {
+		return "", fmt.Errorf("failed to read MFA code: %w", err)
+	}
+	code = strings.TrimSpace(code)
+	if code == "" {
+		return "", fmt.Errorf("MFA code is required")
+	}
+	return code, nil
+}
+
+func selectOktaFactor(factors []oktaFactor) (oktaFactor, error) {
+	var push *oktaFactor
+	for i, factor := range factors {
+		switch factor.FactorType {
+		case oktaFactorTOTP:
+			return factor, nil
+		case oktaFactorPush:
+			if push == nil {
+				push = &factors[i]
+			}
+		}
+	}
+	if push != nil {
+		return *push, nil
+	}
+	types := make([]string, 0, len(factors))
+	for _, factor := range factors {
+		types = append(types, factor.FactorType)
+	}
+	return oktaFactor{}, fmt.Errorf("no supported okta MFA factor found (available: %s); unic currently supports %s and %s",
+		strings.Join(types, ", "), oktaFactorTOTP, oktaFactorPush)
+}
+
+func oktaMFAChallenge(ctx context.Context, client *http.Client, authn oktaAuthnResponse) (string, error) {
+	factor, err := selectOktaFactor(authn.Embedded.Factors)
+	if err != nil {
+		return "", err
+	}
+	if factor.Links.Verify.Href == "" {
+		return "", fmt.Errorf("okta factor %s has no verify link", factor.FactorType)
+	}
+
+	switch factor.FactorType {
+	case oktaFactorTOTP:
+		code, err := promptOktaMFACodeFn(factor)
+		if err != nil {
+			return "", err
+		}
+		resp, err := oktaVerifyFactor(ctx, client, factor.Links.Verify.Href, map[string]string{
+			"stateToken": authn.StateToken,
+			"passCode":   code,
+		})
+		if err != nil {
+			return "", err
+		}
+		if resp.Status != "SUCCESS" || resp.SessionToken == "" {
+			return "", fmt.Errorf("okta MFA verification ended in status %q", resp.Status)
+		}
+		return resp.SessionToken, nil
+
+	case oktaFactorPush:
+		fmt.Fprintln(os.Stderr, "Push notification sent to Okta Verify; waiting for approval...")
+		deadline := time.Now().Add(oktaPushPollTimeout)
+		payload := map[string]string{"stateToken": authn.StateToken}
+		for {
+			resp, err := oktaVerifyFactor(ctx, client, factor.Links.Verify.Href, payload)
+			if err != nil {
+				return "", err
+			}
+			if resp.Status == "SUCCESS" && resp.SessionToken != "" {
+				return resp.SessionToken, nil
+			}
+			if resp.Status == "MFA_CHALLENGE" && resp.FactorResult == "WAITING" {
+				if time.Now().After(deadline) {
+					return "", fmt.Errorf("okta push approval timed out after %s", oktaPushPollTimeout)
+				}
+				select {
+				case <-ctx.Done():
+					return "", ctx.Err()
+				case <-time.After(oktaPushPollInterval):
+				}
+				continue
+			}
+			switch resp.FactorResult {
+			case "REJECTED":
+				return "", fmt.Errorf("okta push was rejected")
+			case "TIMEOUT":
+				return "", fmt.Errorf("okta push timed out")
+			default:
+				return "", fmt.Errorf("okta MFA verification ended in status %q (%s)", resp.Status, resp.FactorResult)
+			}
+		}
+
+	default:
+		return "", fmt.Errorf("unsupported okta MFA factor %q", factor.FactorType)
+	}
+}
+
+func oktaVerifyFactor(ctx context.Context, client *http.Client, href string, payload map[string]string) (oktaAuthnResponse, error) {
+	body, err := json.Marshal(payload)
+	if err != nil {
+		return oktaAuthnResponse{}, err
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, href, bytes.NewReader(body))
+	if err != nil {
+		return oktaAuthnResponse{}, err
+	}
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Accept", "application/json")
+
+	resp, err := client.Do(req)
+	if err != nil {
+		return oktaAuthnResponse{}, fmt.Errorf("okta MFA verification request failed: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode == http.StatusForbidden || resp.StatusCode == http.StatusUnauthorized {
+		return oktaAuthnResponse{}, fmt.Errorf("okta rejected the MFA verification (status %d)", resp.StatusCode)
+	}
+	if resp.StatusCode != http.StatusOK {
+		return oktaAuthnResponse{}, fmt.Errorf("okta MFA verification returned status %d", resp.StatusCode)
+	}
+
+	var verified oktaAuthnResponse
+	if err := json.NewDecoder(resp.Body).Decode(&verified); err != nil {
+		return oktaAuthnResponse{}, fmt.Errorf("failed to parse okta MFA verification response: %w", err)
+	}
+	return verified, nil
+}
+
 // oktaFetchSAMLAssertion loads the app embed link with the one-time session
 // token and extracts the base64 SAML response from the auto-submit form.
 func oktaFetchSAMLAssertion(ctx context.Context, client *http.Client, orgURL, appID, sessionToken string) (string, error) {

internal/auth/okta_test.go

@@ -187,13 +187,110 @@ func TestResolveOktaSAMLSessionReusesCache(t *testing.T) {
 	}
 }
 
-func TestResolveOktaSAMLSessionRejectsMFARequired(t *testing.T) {
+func TestResolveOktaSAMLSessionRejectsUnsupportedFactors(t *testing.T) {
 	server := newOktaTestServer(t, "MFA_REQUIRED")
 	stubOktaSeams(t)
 
 	_, err := ResolveOktaSAMLSession(context.Background(), oktaTestConfig(server.URL))
-	if err == nil || !strings.Contains(err.Error(), "MFA challenge") {
-		t.Fatalf("expected MFA-not-supported error, got %v", err)
+	if err == nil || !strings.Contains(err.Error(), "no supported okta MFA factor") {
+		t.Fatalf("expected unsupported-factor error, got %v", err)
+	}
+}
+
+func newOktaMFATestServer(t *testing.T, factorType string, verifyResponses []string) *httptest.Server {
+	t.Helper()
+	verifyCalls := 0
+	mux := http.NewServeMux()
+	var server *httptest.Server
+	mux.HandleFunc("/api/v1/authn", func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		fmt.Fprintf(w, `{
+			"status": "MFA_REQUIRED",
+			"stateToken": "state-1",
+			"_embedded": {"factors": [{
+				"id": "factor-1",
+				"factorType": %q,
+				"provider": "OKTA",
+				"_links": {"verify": {"href": "%s/api/v1/authn/factors/factor-1/verify"}}
+			}]}
+		}`, factorType, server.URL)
+	})
+	mux.HandleFunc("/api/v1/authn/factors/factor-1/verify", func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		idx := min(verifyCalls, len(verifyResponses)-1)
+		verifyCalls++
+		fmt.Fprint(w, verifyResponses[idx])
+	})
+	mux.HandleFunc("/home/amazon_aws/app123/272", func(w http.ResponseWriter, _ *http.Request) {
+		fmt.Fprintf(w, `<html><input name="SAMLResponse" value="%s"/></html>`, testAssertionB64())
+	})
+	server = httptest.NewServer(mux)
+	t.Cleanup(server.Close)
+	return server
+}
+
+func TestResolveOktaSAMLSessionCompletesTOTPChallenge(t *testing.T) {
+	server := newOktaMFATestServer(t, "token:software:totp", []string{
+		`{"status": "SUCCESS", "sessionToken": "tok-123"}`,
+	})
+	stubOktaSeams(t)
+	origMFA := promptOktaMFACodeFn
+	t.Cleanup(func() { promptOktaMFACodeFn = origMFA })
+	promptOktaMFACodeFn = func(oktaFactor) (string, error) { return "654321", nil }
+
+	session, err := ResolveOktaSAMLSession(context.Background(), oktaTestConfig(server.URL))
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if session.AccessKeyID != "AKIAOKTA" {
+		t.Fatalf("expected exchanged credentials after TOTP, got %+v", session)
+	}
+}
+
+func TestResolveOktaSAMLSessionPollsPushUntilApproved(t *testing.T) {
+	server := newOktaMFATestServer(t, "push", []string{
+		`{"status": "MFA_CHALLENGE", "factorResult": "WAITING"}`,
+		`{"status": "MFA_CHALLENGE", "factorResult": "WAITING"}`,
+		`{"status": "SUCCESS", "sessionToken": "tok-123"}`,
+	})
+	stubOktaSeams(t)
+	origInterval := oktaPushPollInterval
+	t.Cleanup(func() { oktaPushPollInterval = origInterval })
+	oktaPushPollInterval = time.Millisecond
+
+	session, err := ResolveOktaSAMLSession(context.Background(), oktaTestConfig(server.URL))
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if session.AccessKeyID != "AKIAOKTA" {
+		t.Fatalf("expected exchanged credentials after push approval, got %+v", session)
+	}
+}
+
+func TestResolveOktaSAMLSessionFailsOnRejectedPush(t *testing.T) {
+	server := newOktaMFATestServer(t, "push", []string{
+		`{"status": "MFA_CHALLENGE", "factorResult": "REJECTED"}`,
+	})
+	stubOktaSeams(t)
+
+	_, err := ResolveOktaSAMLSession(context.Background(), oktaTestConfig(server.URL))
+	if err == nil || !strings.Contains(err.Error(), "rejected") {
+		t.Fatalf("expected push-rejected error, got %v", err)
+	}
+}
+
+func TestResolveOktaSAMLSessionFailsOnBadTOTPCode(t *testing.T) {
+	server := newOktaMFATestServer(t, "token:software:totp", []string{
+		`{"status": "MFA_CHALLENGE", "factorResult": "REJECTED"}`,
+	})
+	stubOktaSeams(t)
+	origMFA := promptOktaMFACodeFn
+	t.Cleanup(func() { promptOktaMFACodeFn = origMFA })
+	promptOktaMFACodeFn = func(oktaFactor) (string, error) { return "000000", nil }
+
+	_, err := ResolveOktaSAMLSession(context.Background(), oktaTestConfig(server.URL))
+	if err == nil || !strings.Contains(err.Error(), "MFA verification ended") {
+		t.Fatalf("expected TOTP verification failure, got %v", err)
 	}
 }