A stateful CLI tool for getting access token from an OIDC provider

oidc_client.clj

#!/usr/bin/env bb
(ns oidc-client
  "Get end-user access token from an OIDC provider, caching the access/refresh token in an encrypted file. Code in https://babashka.org
   Usage:
   
       /path/to/oidc_client.clj
   
   When there are no cached, valid tokens, it will open a browser with the OIDC provider login URL and start a HTTPS-enabled
   callback server on localhost. Make sure that the resulting redirect_uri is registered with your provider.
  
   NOTE: After login you will be redirected to https://<localhost> that uses a self-signed certificate and will need to tell your
         browser to trust the site despite all the danger warnings :).
  
  TIP: Use https://github.com/FiloSottile/mkcert/ to create a trusted certificate and mount it
       to the proxy - see https://github.com/dweomer/dockerfiles-stunnel
   
   Prerequisities: babashka (e.g. v0.3.4), openssl, config file with credentials."
  (:require
   [babashka.fs :as fs]
   [babashka.process :as proc :refer [process]]
   [cheshire.core :as json]
   [clojure.edn :as edn]
   [clojure.java.shell :as sh]
   [clojure.string :as str]
   [org.httpkit.client :as http]
   [org.httpkit.server :as server]))

;;--------------------------------------------------------------------------------------------------------- vars
(def encryption-psw "TODO: put some random string here")

(def oidc-config
  {:client-id "<your oidc client id>"
   :client-secret "<your oidc client id>"
    assert)
   :token-url "https://oidc-provider.example.com/openid-connect/token"
   :login-url "https://oidc-provider.example.com/openid-connect/auth?client_id=<your oidc client id>&prompt=login&redirect_uri=https%3A%2F%2Flocalhost%3A8080%2Fapi%2Foauth2%2Fcallback&response_type=code&scope=openid"
   :callback-url {:host "localhost" :port 8080 :path "/api/oauth2/callback"}})

(defn callback-url [oidc-config]
  (let [{:keys [host port path]} (:callback-url oidc-config)]
    (str "https://" host \: port path)))

;;--------------------------------------------------------------------------------------------------------- fs cache

(defn encrypt [plaintext]
  (let [{:keys [exit out err]}
        (sh/sh "openssl" "enc" "-e" "-des3" "-base64" "-pass" (str "pass:" encryption-psw) "-pbkdf2"
               :in plaintext)]
    (when-not (zero? exit)
      (throw (ex-info (str "encryption failed with status " exit " and msg " err) {})))
    out))

(defn decrypt [cyphertext]
  (let [{:keys [exit out err]}
        (sh/sh "openssl" "enc" "-d" "-des3" "-base64" "-pass" (str "pass:" encryption-psw) "-pbkdf2" 
         :in cyphertext)]
    (when-not (zero? exit)
      (throw (ex-info (str "decryption failed with status " exit " and msg " err) {})))
    out))

;;---------------------------------------------------------------------------------------------------------

(defn cache-dir []
  (let [osx-cache  (str (System/getProperty "user.home") "/Library/Caches") #_osx
        cache-root (or (System/getenv "XDG_CACHE_HOME") #_linux
                       (System/getenv "LOCALAPPDATA") #_windows
                       (when (fs/directory? osx-cache) osx-cache))
        dir        (fs/file cache-root "oidc_client_bb")]
    (when-not (fs/directory? cache-root)
      (throw (ex-info (str "Guessed cache data dir '" cache-root "' does not exist") {})))
    (when-not (fs/directory? dir)
      (fs/create-dir dir)
      (fs/set-posix-file-permissions dir "rwx------"))
    dir))

(defn cache-data [file-name data {:keys [encrypt?]}]
  (spit (fs/file (cache-dir) file-name) 
        (cond-> (pr-str data)
                encrypt? (encrypt))))

(defn read-cached-data [file-name {:keys [decrypt?]}]
  (let [file (fs/file (cache-dir) file-name)]
    (when (fs/readable? file)
      (-> (slurp file)
          (cond->
           decrypt? (decrypt))
          (edn/read-string)))))

;;--------------------------------------------------------------------------------------------------------- http server

(defn check-port-free [port]
  (try
    (with-open [s (java.net.ServerSocket. port)] true)
    (catch Exception e ; java.net.BindException not known to babashka
      (throw (ex-info (str "Cannot proceed, the port " port " is already occupied") {:port port, :e e})))))

(defn free-port []
  (with-open [s (java.net.ServerSocket. 0)]
    (.getLocalPort s)))

(defn start-https-proxy
  "Start the HTTPS process and return a Deref that will destroy it"
  [upstream-port]
  (check-port-free 8080)
  (let [https-port (-> oidc-config :callback-url :port)
        upstream-var (str "STUNNEL_CONNECT=host.docker.internal:" upstream-port)
        ;; BEWARE: host.docker.internal works on OSX, likely diff. on Win
        proc (proc/$ docker run --rm --name bb-oidc-stunnel
                     -e STUNNEL_SERVICE=bb-oidc-client
                     -e ~(str "STUNNEL_ACCEPT=" https-port)
                     -e ~upstream-var
                     -p ~(str https-port \: https-port)
                     "dweomer/stunnel@sha256:3601510afa54b2dc1378b4d7cd25c4bd96180201ad254fb76b864fcc4e0f5dfd")]
    (babashka.wait/wait-for-port "localhost" https-port {:timeout 3000 :pause 1000})
    (when-not (-> proc :proc .isAlive)
      (throw (ex-info (format "Starting a https proxy at %d failed with status %d and err: %s"
                              https-port
                              (:exit @proc)
                              (slurp (:err @proc)))
                      {:proc proc})))
    (delay (proc/destroy proc))))

(defn format-tokens-response [{:keys [access_token expires_in refresh_expires_in refresh_token]}]
  (let [now-ms (System/currentTimeMillis)
        ->expires-at #(+ now-ms (* % 1000))]
    {:access  {:token access_token,  :expires-at-ms (->expires-at expires_in)}
     :refresh {:token refresh_token, :expires-at-ms (->expires-at refresh_expires_in)}}))

(defn fetch-identity-token
  ([operation code-or-token]
   {:pre [(#{:login :refresh} operation) (string? code-or-token)]}
   @(http/post
     (:token-url oidc-config)
     {:timeout 10000 :connect-timeout 5000
      :proxy-url (System/getenv "https_proxy")
      :headers {"accept" "application/json"}
      :basic-auth ((juxt :client-id :client-secret) oidc-config)
      :form-params (if (= operation :login)
                     {:grant_type "authorization_code"
                      :redirect_uri (callback-url oidc-config) ; required for some reason
                      :code code-or-token}
                     {:grant_type "refresh_token"
                      :refresh_token code-or-token})}
     (fn [{:keys [opts status body headers error] :as resp}]
       (cond
         error
         (ex-info "Network request failed" {} error)

         (>= status 300)
         (ex-info (str "Got error status " status ": " body) {:status status, :body body})

         :else
         ;; res = {:access_token, :expires_in [s], :refresh_expires_in, :refresh_token, :session_state, ...}
         (-> body (json/parse-string true) format-tokens-response))))))

(defn parse-query-string [query]
  (->> (str/split query #"&")
       (mapcat #(str/split % #"="))
       (apply hash-map)))

(defn redeem-code->tokens [query-string]
  (let [code (-> (parse-query-string query-string)
                 (get "code"))
        res  (fetch-identity-token :login code)]
    (when (instance? java.lang.Exception res)
      (throw res))
    res))

(defn handler [tokensp {:keys [uri query-string request-method] :as req}]
  (cond
    (and (= uri (-> oidc-config :callback-url :path))
         query-string)
    (try
      (->> (redeem-code->tokens query-string)
           (deliver tokensp))
      {:body "Tokens registered, you can now close this window"}
      (catch java.lang.Exception e
        (deliver tokensp nil)
        {:status 500, :body (str e)}))

    :else {:status 404 :body (str "Unknown page " uri)}))

(defn oidc-login []
  (let [tokensp (promise)
        backend-port (free-port)
        stop-proxy (start-https-proxy backend-port)
        srv (server/run-server (partial handler tokensp) {:port backend-port, :legacy-return-value? false})
        _ (clojure.java.browse/browse-url (:login-url oidc-config))
        tokens @tokensp]
    @stop-proxy
    (server/server-stop! srv)
    tokens))

;;--------------------------------------------------------------------------------------------------------- main

(defn token-valid? 
  ([tokens] (token-valid? tokens nil))
  ([{:keys [expires-at-ms]} min-ttl-ms]
   (and expires-at-ms
        (> expires-at-ms
           (+ (System/currentTimeMillis) (or min-ttl-ms 1000))))))

(defn store-tokens [tokens]
  (cache-data "tokens.enc" tokens {:encrypt? true})
  (binding [*out* *err*] (println "[log] auth-token: stored new tokens, valid until" (java.util.Date. (-> tokens :refresh :expires-at-ms))))
  tokens)

(defn get-auth-token []
  (let [tokens         (read-cached-data "tokens.enc" {:decrypt? true})
        valid-acces?   (token-valid? (:access tokens) (* 10 60 1000))
        valid-refresh? (token-valid? (:refresh tokens))]
    (binding [*out* *err*] (println "[log] auth-token:" 
                                    (cond
                                      valid-acces? (str "returning the cached access token, valid until " (java.util.Date. (-> tokens :access :expires-at-ms)))
                                      valid-refresh? "going to refresh the tokens..."
                                      :else "no cached valid token, logging in...")))
    (cond valid-acces?
          (-> tokens :access :token println)

          valid-refresh?
          (-> (doto (fetch-identity-token :refresh (-> tokens :refresh :token))
                (store-tokens))
              :access :token
              println)

          :else
          (-> (doto (oidc-login)
                (store-tokens))
              :access :token
              println))))

(when (= *file* (System/getProperty "babashka.file"))
  (get-auth-token))