metabase-enterprise.sso.integrations.jwt

(ns metabase-enterprise.sso.integrations.jwt
  (:require
   [buddy.sign.jwt :as jwt]
   [clojure.string :as str]
   [java-time.api :as t]
   [metabase-enterprise.sso.api.interface :as sso.i]
   [metabase-enterprise.sso.integrations.sso-settings :as sso-settings]
   [metabase-enterprise.sso.integrations.sso-utils :as sso-utils]
   [metabase.embed.settings :as embed.settings]
   [metabase.premium-features.core :as premium-features]
   [metabase.request.core :as request]
   [metabase.session.models.session :as session]
   [metabase.sso.core :as sso]
   [metabase.util.i18n :refer [tru]]
   [ring.util.response :as response])
  (:import
   (java.net URLEncoder)))

(set! *warn-on-reflection* true)

(defn fetch-or-create-user!
  [first-name last-name email user-attributes]
  (when-not (sso-settings/jwt-enabled)
    (throw
     (IllegalArgumentException.
      (str (tru "Can''t create new JWT user when JWT is not configured")))))
  (let [user {:first_name       first-name
              :last_name        last-name
              :email            email
              :sso_source       :jwt
              :login_attributes user-attributes}]
    (or (sso-utils/fetch-and-update-login-attributes! user)
        (sso-utils/check-user-provisioning :jwt)
        (sso-utils/create-new-sso-user! user))))

(def ^:private ^{:arglists '([])} jwt-attribute-email
  (comp keyword sso-settings/jwt-attribute-email))

(def ^:private ^{:arglists '([])} jwt-attribute-firstname
  (comp keyword sso-settings/jwt-attribute-firstname))

(def ^:private ^{:arglists '([])} jwt-attribute-lastname
  (comp keyword sso-settings/jwt-attribute-lastname))

(def ^:private ^{:arglists '([])} jwt-attribute-groups
  (comp keyword sso-settings/jwt-attribute-groups))

(def ^:private registered-claims
  [:iss :iat :sub :aud :exp :nbf :jti])

(defn- jwt-data->login-attributes [jwt-data]
  (apply dissoc
         jwt-data
         (jwt-attribute-email)
         (jwt-attribute-firstname)
         (jwt-attribute-lastname)
         registered-claims))

(def ^:private ^:const three-minutes-in-seconds 180)

(defn- group-names->ids
  [group-names]
  (set
   (mapcat (sso-settings/jwt-group-mappings)
           (map keyword group-names))))

(defn- all-mapped-group-ids
  []
  (-> (sso-settings/jwt-group-mappings)
      vals
      flatten
      set))

(defn- sync-groups!
  [user jwt-data]
  (when (sso-settings/jwt-group-sync)
    (when-let [groups-attribute (jwt-attribute-groups)]
      (when-let [group-names (get jwt-data groups-attribute)]
        (sso/sync-group-memberships! user
                                     (group-names->ids group-names)
                                     (all-mapped-group-ids))))))

(defn- session-data
  [jwt {{redirect :return_to} :params, :as request}]
  (let [redirect-url (or redirect (URLEncoder/encode "/"))]
    (sso-utils/check-sso-redirect redirect-url)
    (let [jwt-data     (try
                         (jwt/unsign jwt (sso-settings/jwt-shared-secret)
                                     {:max-age three-minutes-in-seconds})
                         (catch Throwable e
                           (throw
                            (ex-info (ex-message e)
                                     {:status      "error-jwt-bad-unsigning"
                                      :status-code 401}))))
          login-attrs  (jwt-data->login-attributes jwt-data)
          email        (get jwt-data (jwt-attribute-email))
          first-name   (get jwt-data (jwt-attribute-firstname))
          last-name    (get jwt-data (jwt-attribute-lastname))
          user         (fetch-or-create-user! first-name last-name email login-attrs)
          session      (session/create-session! :sso user (request/device-info request))]
      (sync-groups! user jwt-data)
      {:session session, :redirect-url redirect-url, :jwt-data jwt-data})))

(defn- check-jwt-enabled []
  (when-not (sso-settings/jwt-configured)
    (throw
     (ex-info (tru "JWT SSO has not been configured")
              {:status      "error-sso-jwt-not-configured"
               :status-code 402})))
  (when-not (sso-settings/jwt-enabled)
    (throw
     (ex-info (tru "JWT SSO has not been enabled")
              {:status      "error-sso-jwt-disabled"
               :status-code 402})))
  true)

(defn ^:private generate-response-token
  [session jwt-data]
  (if-not (embed.settings/enable-embedding-sdk)
    (throw
     (ex-info (tru "SDK Embedding is disabled. Enable it in the Embedding settings.")
              {:status      "error-embedding-sdk-disabled"
               :status-code 402}))
    (response/response
     {:status :ok
      :id     (:id session)
      :exp    (:exp jwt-data)
      :iat    (:iat jwt-data)})))

(defn ^:private redirect-to-idp
  [idp redirect]
  (let [return-to-param (if (str/includes? idp "?") "&return_to=" "?return_to=")]
    (response/redirect
     (str idp
          (when redirect
            (str return-to-param redirect))))))

(defn ^:private handle-jwt-authentication
  [{:keys [session redirect-url jwt-data]} token request]
  (if token
    (generate-response-token session jwt-data)
    (request/set-session-cookies request (response/redirect redirect-url) session (t/zoned-date-time (t/zone-id "GMT")))))

(defmethod sso.i/sso-get :jwt
  [{{:keys [jwt redirect token] :or {token nil}} :params, :as request}]
  (premium-features/assert-has-feature :sso-jwt (tru "JWT-based authentication"))
  (check-jwt-enabled)
  (if jwt
    (handle-jwt-authentication (session-data jwt request) token request)
    (redirect-to-idp (sso-settings/jwt-identity-provider-uri) redirect)))

(defmethod sso.i/sso-post :jwt
  [_]
  (throw
   (ex-info (tru "POST not valid for JWT SSO requests")
            {:status "error-post-jwt-not-valid" :status-code 501})))