AWS Cognito

(ns cognito

  {:nextjournal.clerk/toc true}

  (:require [buddy.core.hash :as hash]

            [buddy.sign.jwt :as jwt]

            [buddy.core.codecs :as codecs

              :refer [to-bytes bytes->str 

                      bytes->hex hex->bytes

                      bytes->b64 b64->bytes]]

            [buddy.core.keys :refer [jwk->public-key]]

            [clojure.string :as str]

            [clojure.set :refer [rename-keys]]

            [clojure.pprint :refer [pprint]]

            [java-time.api :as jt]

            [cheshire.core :as json]

            [clj-http.client :as client])

  (:import [java.math BigInteger]

           [java.security SecureRandom]

           [javax.crypto Mac]

           [javax.crypto.spec SecretKeySpec]))

(extend-protocol codecs/IByteArray

  BigInteger

  (-to-bytes [data] (.toByteArray data)))

(defn pad-hex

  "BigInteger here is unsigned but the hex string could result in a negative number

  unless padded with some 0's"

  [^String hex]

  (cond

    (-> (count hex) (mod 2) (= 1)) (str "0" hex)

    (->> (first hex) (str) (str/includes? "89ABCDEFabcdef")) (str "00" hex)

    :else hex))

(defn big-integer-to-hex

  [^BigInteger big-integer]

  (bytes->hex (to-bytes big-integer)))

(def user-pool-config-public

  "allows USER_PASSWORD_AUTH flow (no client_secret)"

  (merge {:user-pool-id "us-east-2_INzuX7Ms0"

          :client-id "mcedtgsalss1qln1q4gm84r3s"}

    (json/parse-string (slurp "./.ov_password.secret") true)))

(def user-pool-config

  "requires USER_SRP_AUTH flow (w/ client_secret)"

  (merge {:user-pool-id "us-east-2_INzuX7Ms0"

          :client-id "hubekaq48puviiu82akhrvu0m"

          :client-secret "1dgn40avs48ek0sompk784a1fgv5itnnkstpk7ed1mbs4oilfala"}

    (json/parse-string (slurp "./.ov_password.secret") true)))

""

(def ^:dynamic *known-token-keys* (atom {}))

(defn get-openid-config

  [user-pool-id]

  (let [region (first (str/split user-pool-id #"_"))]

    (-> (format "https://cognito-idp.%s.amazonaws.com/%s/.well-known/openid-configuration"

                region user-pool-id)

        (client/get {:as :json})

      :body)))

(defn add-user-pool-keys

  "map of :kid (key id from the token header) to an actual public key"

  [user-pool-id]

  (-> (get-openid-config user-pool-id)

      :jwks_uri

      ;; now we have the uri to get the list of authorized keys so go get them

      (client/get {:as :json})

      :body

      :keys

      (->> (into {} (map (juxt :kid jwk->public-key)))

        (swap! *known-token-keys* merge))))

(defn public-key-from-token [token]

 ;mostly borrowed from jws/parse-header-

  (let [head-data (-> token 

                      (str/split #"\.")

                      (first)

                      (to-bytes)

                      (b64->bytes)

                      (bytes->str)

                      (json/parse-string true))

        public-key (@*known-token-keys* (:kid head-data))

        alg (keyword (str/lower-case (:alg  head-data)))]

    (when public-key

      [public-key {:alg alg}])))

(-> (:user-pool-id user-pool-config)

             (add-user-pool-keys))

(defn init-auth-parameters

  "Does not use SRP and does not have client-secret set.

   Therefore, we can use USER_PASSWORD_AUTH flow"

  [{:keys [client-id USERNAME PASSWORD] :as init-params}]

  (assoc init-params :init-auth-params {"AuthFlow" "USER_PASSWORD_AUTH"

                                        "ClientId" client-id   

                                        "AuthParameters" {"USERNAME" USERNAME

                                                          "PASSWORD" PASSWORD}}))

(def init-auth-headers

  {

   "X-Amz-Target" "AWSCognitoIdentityProviderService.InitiateAuth"

   "Content-Type" "application/x-amz-json-1.1" ;; WARNING: do now override with {:as :json}

  })

(defn get-tokens

  [{:keys [user-pool-id init-auth-params]}]

  (-> (format "https://cognito-idp.%s.amazonaws.com/login" (first (str/split user-pool-id #"_" 2)))

      (client/post {:headers init-auth-headers

                    :body (json/generate-string init-auth-params)})

      :body

      (json/parse-string true)

      :AuthenticationResult

      ; the token and login endpoints have different key names so lets fix that up

      (rename-keys {:RefreshToken :refresh_token

                    :AccessToken :access_token

                    :IdToken :id_token

                    :TokenType :token_type

                    :ExpiresIn :expires_in})))

(defn unsign-cognito [token]

  (apply jwt/unsign token (public-key-from-token token)))

(->  user-pool-config-public

     init-auth-parameters

     get-tokens

     :id_token

     unsign-cognito

     pprint)

(defn sign-to-bytes

  "Returns the byte signature of a string with a given key.

  Uses optional {:algorithm xxx} name which deafults to \"HmacSHA256\"."

  [key string & {:keys [algorithm] :or {algorithm "HmacSHA256"}}]

  (let [mac (Mac/getInstance algorithm)

        secretKey (SecretKeySpec. (to-bytes key) (.getAlgorithm mac))]

    (-> (doto mac

          (.init secretKey)

          (.update (to-bytes string)))

        .doFinal)))

(defn sign-to-b64

  "Returns the HMAC SHA256 b64 string signature from a key-string pair."

  [key string]

  (-> (sign-to-bytes key string)

      (bytes->b64)

      (bytes->str)))

(defn get-secret-hash

  [username client-id client-secret]

  (->> (str username client-id)

       (sign-to-b64 client-secret)))

;; these are just two constants

(def g (BigInteger. "2"))

(def big-n

  (-> (str "FFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD1"

           "29024E088A67CC74020BBEA63B139B22514A08798E3404DD"

           "EF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245"

           "E485B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7ED"

           "EE386BFB5A899FA5AE9F24117C4B1FE649286651ECE45B3D"

           "C2007CB8A163BF0598DA48361C55D39A69163FA8FD24CF5F"

           "83655D23DCA3AD961C62F356208552BB9ED529077096966D"

           "670C354E4ABC9804F1746C08CA18217C32905E462E36CE3B"

           "E39E772C180E86039B2783A2EC07A28FB5C55DF06F4C52C9"

           "DE2BCBF6955817183995497CEA956AE515D2261898FA0510"

           "15728E5A8AAAC42DAD33170D04507A33A85521ABDF1CBA64"

           "ECFB850458DBEF0A8AEA71575D060C7DB3970F85A6E1E4C7"

           "ABF5AE8CDB0933D71E8C94E04A25619DCEE3D2261AD2EE6B"

           "F12FFA06D98A0864D87602733EC86A64521F2B18177B200C"

           "BBE117577A615D6C770988C0BAD946E208E24FA074E5AB31"

           "43DB5BFCE0FD108E4B82D120A93AD2CAFFFFFFFFFFFFFFFF")

    (BigInteger. 16)))

;; k is calculated but it is calculated from constants

;; so let's just skip that and use the resluting BigInteger

;; from python: hex_to_long(hex_hash('00' + n_hex + '0' + g_hex))

;; from clojure (BigInteger. (hash/sha256 (hex->bytes (apply str (map #(bytes->hex (.toByteArray %)) [big-n g])))))

(def k

  (BigInteger. "37772559067617796309459009502931177628717927509759535181635788491848250400486"))

(def DERIVED-KEY "Caldera Derived Key")

(defn modPow

  "modPow that is easier to use as arity 1 for the common case"

  ([small-x] (modPow g small-x big-n))

  ([g small-x n] (.modPow g small-x big-n)))

(defn init-srp-parameters

  [{:keys [client-id client-secret USERNAME small-a big-a] :as init-params}]

  (let [params (cond-> init-params

                 (and small-a (nil? big-a)) (assoc :big-a (modPow small-a))

                 (nil? small-a) (#(let [new-small-a (BigInteger. (* 128 8) (SecureRandom.))]

                                   (assoc % :small-a new-small-a

                                            :big-a   (modPow new-small-a)))))

        auth-params (cond-> {"ClientId" client-id

                             "AuthFlow" "USER_SRP_AUTH"

                             "AuthParameters" {"USERNAME" USERNAME

                                               "SRP_A" (bytes->hex (to-bytes (:big-a params)))}}

                      client-secret (assoc-in ["AuthParameters" "SECRET_HASH"]

                                              (get-secret-hash USERNAME client-id client-secret)))]

     (assoc params :init-auth-params auth-params)))

(defn get-srp-challenge

  [{:keys [user-pool-id init-auth-params] :as init-config}]

  (let [challenge (-> (format "https://cognito-idp.%s.amazonaws.com/login" (first (str/split user-pool-id #"_" 2)))

                      (client/post {:headers init-auth-headers

                                    :body (json/generate-string init-auth-params)}))

        request-id (-> challenge

                       :headers

                       (get "x-amzn-RequestId"))

        auth-challenge (-> challenge                      

                           :body

                           (json/parse-string true))]

       (assoc init-config

              :auth-challenge auth-challenge

              :request-id request-id)))

(-> (init-srp-parameters user-pool-config)

    (get-srp-challenge)

    (assoc :PASSWORD "<no-peaking>"))

(defn calculate-u

  [big-a big-b]

  ;; TODO: sort out "0" padding at the front

  (-> (str (big-integer-to-hex big-a) (big-integer-to-hex big-b))

      (hex->bytes)

      (hash/sha256)

      (bytes->hex)

      (BigInteger. 16)))

(defn compute-hkdf [^bytes s  ^bytes u]

  (let [mac (Mac/getInstance "HmacSHA256")

        prk (-> (doto mac

                  (.init (SecretKeySpec. u "SHA256"))

                  (.update s))

              (.doFinal))

        info-bits (to-bytes (str "Caldera Derived Key" (char 1)))

        hmac-hash (-> (doto mac

                        (.reset)

                        (.init (SecretKeySpec. prk "SHA256"))

                        (.update info-bits))

                    (.doFinal))]

    (byte-array (take 16 hmac-hash))))

(defn get-password-authentication-key

  [pool-id user-name password small-a big-a big-b salt]

  (let [u-value (calculate-u big-a big-b)

        pool (second (str/split pool-id #"_" 2))

        username-password (-> (format "%s%s:%s" pool user-name password)

                            (hash/sha256)

                            (bytes->hex))

        x-value (-> (pad-hex salt)

                  (str username-password)

                  (hex->bytes)

                  (hash/sha256)

                  (bytes->hex) ; TODO: do we need another (pad-hex) here?

                  (BigInteger. 16))

        g-mod-pow-xn (.modPow g x-value big-n)

        int-val2 (->> (.modPow g x-value big-n)

                   (.multiply k)

                   (.subtract big-b))

        s-value (.modPow

                  int-val2

                  (-> (.multiply u-value x-value)

                    (.add small-a))

                  big-n)]

    (compute-hkdf

      (hex->bytes (pad-hex (bytes->hex (.toByteArray s-value))))

      (hex->bytes (pad-hex (bytes->hex (.toByteArray u-value)))))))

(defn- get-timestamp

  []

  (jt/format "EEE MMM d HH:mm:ss z yyyy" (jt/zoned-date-time (jt/instant) "UTC")))

(defn process-challenge

  "process a challenge"

  [{:keys [user-pool-id small-a big-a client-id client-secret PASSWORD] :as config}]

  (let [{salt-hex :SALT

         big-b-hex :SRP_B

         secret-block-b64 :SECRET_BLOCK

         user-id-for-srp :USER_ID_FOR_SRP

         challenge-username :USERNAME} (get-in config [:auth-challenge :ChallengeParameters])

        big-b (BigInteger. big-b-hex 16)

        u-value (calculate-u big-a big-b)

        secret-block-bytes (-> secret-block-b64

                               (to-bytes)

                               (b64->bytes))

        timestamp (get-timestamp)

        pool-name (second (str/split user-pool-id #"_" 2)) ;; region_poolname => poolname

        message  (->> [pool-name challenge-username secret-block-bytes timestamp]

                      (mapcat to-bytes)

                      (byte-array))

        hkdf   (get-password-authentication-key

                user-pool-id 

                user-id-for-srp

                PASSWORD

                small-a 

                big-a 

                big-b

                salt-hex)

        response-params (cond-> {"TIMESTAMP" timestamp

                                 "USERNAME" challenge-username

                                 "PASSWORD_CLAIM_SECRET_BLOCK" secret-block-b64

                                 "PASSWORD_CLAIM_SIGNATURE" (sign-to-b64 hkdf message)}

                          client-secret (assoc "SECRET_HASH" (get-secret-hash

                                                                  user-id-for-srp

                                                                  client-id

                                                                  client-secret)))]    

    (assoc config :challenge-response {"ClientId" client-id

                                       "ChallengeName" "PASSWORD_VERIFIER"

                                       "ChallengeResponses" response-params})))

(def challenge-response-headers

  "partial header for a challenge response.

  This needs the `x-amzn-RequestId` to be added before using"

  {"X-Amz-Target" "AWSCognitoIdentityProviderService.RespondToAuthChallenge",

   "Content-Type" "application/x-amz-json-1.1"})

(defn submit-challenge-response

  [{:keys [user-pool-id challenge-response request-id]}]

 (try

   (-> (client/post "https://cognito-idp.us-east-2.amazonaws.com/"

                    {:headers (assoc challenge-response-headers

                                     "x-amzn-RequestId" request-id)

                     :body (json/generate-string challenge-response)})

       :body

       (json/parse-string true)

       :AuthenticationResult

       ; the token and login endpoints have different key names so lets fix that up

       (rename-keys {:RefreshToken :refresh_token

                     :AccessToken :access_token

                     :IdToken :id_token

                     :TokenType :token_type

                     :ExpiresIn :expires_in}))

   (catch Exception e (-> e (Throwable->map)

                          :data

                          ((juxt :status #(json/parse-string (:body %))))))))

(-> user-pool-config

    (init-srp-parameters)

    (get-srp-challenge)

    (process-challenge)

    (submit-challenge-response))

;; the "public" pool also accepts the SRP_AUTH flow

(-> user-pool-config-public

    (init-srp-parameters)

    (get-srp-challenge)

    (process-challenge)

    (submit-challenge-response))