#lang racket/base

(require crypto
         crypto/argon2
         net/private/ip
         racket/list
         racket/random
         racket/string
         web-server/http
         web-server/http/cookie-parse)

(provide make-password-hash
         password-hash-valid?
         make-auth-manager
         auth-manager?
         auth-enabled?
         auth-request-local?
         auth-request-user
         auth-login!
         auth-logout!
         auth-session-cookie
         auth-expired-cookie)

(struct ip-network (address prefix) #:transparent)
(struct session (username [last-seen #:mutable]) #:transparent)
(struct failures ([attempts #:mutable] [started #:mutable]) #:transparent)
(struct auth-manager
  (users local-networks trusted-proxies session-seconds sessions failed lock)
  #:transparent)

(define password-kdf
  (or (get-kdf 'argon2id argon2-factory)
      (error 'rkt-web-player/users "Argon2id is unavailable")))

(define password-parameters
  '((m 19456) (t 2) (p 1)))

;; Used to make an unknown username take the same expensive verification path.
(define dummy-password-hash
  "$argon2id$v=19$m=19456,t=2,p=1$WrJi0t7NsbD3adX8kxrT/g$yCIf2Ork8G8PRdIcA0bEIYIdMwzunrTou/BKcM4cO/0")

(define session-cookie-name "rkt-web-player-session")
(define failure-window-seconds 300)
(define maximum-failures 5)

(define (make-password-hash password)
  (unless (and (string? password)
               (>= (string-length password) 12))
    (raise-argument-error
     'make-password-hash
     "string containing at least 12 characters"
     password))
  (pwhash password-kdf
          (string->bytes/utf-8 password)
          password-parameters))

(define (password-hash-valid? password encoded)
  (and (string? password)
       (string? encoded)
       (with-handlers ((exn:fail? (lambda (_) #f)))
         (pwhash-verify password-kdf
                        (string->bytes/utf-8 password)
                        encoded))))

(define (normal-ip-bytes value)
  (define raw
    (ip-address->bytes (make-ip-address value)))
  ;; Normalize IPv4-mapped IPv6 addresses to four bytes.
  (if (and (= (bytes-length raw) 16)
           (for/and ((index (in-range 10)))
             (zero? (bytes-ref raw index)))
           (= (bytes-ref raw 10) #xff)
           (= (bytes-ref raw 11) #xff))
      (subbytes raw 12)
      raw))

(define (parse-network value)
  (define parts (string-split (string-trim value) "/"))
  (unless (member (length parts) '(1 2))
    (raise-argument-error 'make-auth-manager "IP address or CIDR network" value))
  (define address
    (with-handlers ((exn:fail?
                     (lambda (_)
                       (raise-argument-error
                        'make-auth-manager
                        "IP address or CIDR network"
                        value))))
      (normal-ip-bytes (car parts))))
  (define maximum (* 8 (bytes-length address)))
  (define prefix
    (if (= (length parts) 2)
        (string->number (cadr parts))
        maximum))
  (unless (and (exact-nonnegative-integer? prefix)
               (<= prefix maximum))
    (raise-argument-error 'make-auth-manager "IP address or CIDR network" value))
  (ip-network address prefix))

(define (network-contains? network address-string)
  (with-handlers ((exn:fail? (lambda (_) #f)))
    (define candidate (normal-ip-bytes address-string))
    (define expected (ip-network-address network))
    (and (= (bytes-length candidate) (bytes-length expected))
         (let-values (((whole remainder)
                       (quotient/remainder (ip-network-prefix network) 8)))
           (and (for/and ((index (in-range whole)))
                  (= (bytes-ref candidate index)
                     (bytes-ref expected index)))
                (or (zero? remainder)
                    (let ((mask
                           (bitwise-and #xff
                                        (arithmetic-shift #xff (- remainder 8)))))
                      (= (bitwise-and (bytes-ref candidate whole) mask)
                         (bitwise-and (bytes-ref expected whole) mask)))))))))

(define (header-string request name)
  (let ((value (headers-assq* name (request-headers/raw request))))
    (and value
         (bytes->string/utf-8 (header-value value)))))

(define (trusted-proxy? manager address)
  (ormap (lambda (network) (network-contains? network address))
         (auth-manager-trusted-proxies manager)))

(define (request-address manager request)
  (define peer (request-client-ip request))
  (define forwarded
    (and (trusted-proxy? manager peer)
         (header-string request #"X-Forwarded-For")))
  (if forwarded
      ;; A trusted reverse proxy appends the address it observed. Earlier
      ;; values can have been supplied by the untrusted client.
      (string-trim (last (string-split forwarded ",")))
      peer))

(define (auth-enabled? manager)
  (positive? (hash-count (auth-manager-users manager))))

(define (auth-request-local? manager request)
  (ormap (lambda (network)
           (network-contains? network (request-address manager request)))
         (auth-manager-local-networks manager)))

(define (request-session-token request)
  (for/or ((cookie (in-list (request-cookies request))))
    (and (string=? (client-cookie-name cookie) session-cookie-name)
         (client-cookie-value cookie))))

(define (prune-sessions! manager now)
  (for ((token (in-list (hash-keys (auth-manager-sessions manager)))))
    (let ((value (hash-ref (auth-manager-sessions manager) token)))
      (when (> (- now (session-last-seen value))
               (auth-manager-session-seconds manager))
        (hash-remove! (auth-manager-sessions manager) token)))))

(define (auth-request-user manager request)
  (cond
    ((not (auth-enabled? manager)) "anonymous")
    ((auth-request-local? manager request) "local")
    (else
     (let ((token (request-session-token request))
           (now (current-seconds)))
       (and token
            (call-with-semaphore
             (auth-manager-lock manager)
             (lambda ()
               (prune-sessions! manager now)
               (let ((value (hash-ref (auth-manager-sessions manager)
                                      token
                                      #f)))
                 (and value
                      (begin
                        (set-session-last-seen! value now)
                        (session-username value)))))))))))

(define (failure-blocked? manager address now)
  (define value (hash-ref (auth-manager-failed manager) address #f))
  (and value
       (if (> (- now (failures-started value)) failure-window-seconds)
           (begin
             (hash-remove! (auth-manager-failed manager) address)
             #f)
           (>= (failures-attempts value) maximum-failures))))

(define (record-failure! manager address now)
  (define value (hash-ref (auth-manager-failed manager) address #f))
  (if (and value
           (<= (- now (failures-started value)) failure-window-seconds))
      (set-failures-attempts! value (+ 1 (failures-attempts value)))
      (hash-set! (auth-manager-failed manager)
                 address
                 (failures 1 now))))

;; Returns a new token, #f for invalid credentials, or 'rate-limited.
(define (auth-login! manager request username password)
  (define address (request-address manager request))
  (define now (current-seconds))
  (define normalized (string-downcase (string-trim username)))
  (call-with-semaphore
   (auth-manager-lock manager)
   (lambda ()
     (if (failure-blocked? manager address now)
         'rate-limited
         (let* ((stored (hash-ref (auth-manager-users manager)
                                  normalized
                                  #f))
                (valid?
                 (password-hash-valid?
                  password
                  (or stored dummy-password-hash))))
           (if (and stored valid?)
               (let ((token
                      (bytes->hex-string (crypto-random-bytes 32))))
                 (hash-remove! (auth-manager-failed manager) address)
                 (prune-sessions! manager now)
                 (hash-set! (auth-manager-sessions manager)
                            token
                            (session normalized now))
                 token)
               (begin
                 (record-failure! manager address now)
                 #f)))))))

(define (auth-logout! manager request)
  (let ((token (request-session-token request)))
    (when token
      (call-with-semaphore
       (auth-manager-lock manager)
       (lambda ()
         (hash-remove! (auth-manager-sessions manager) token))))))

(define (auth-session-cookie manager token)
  (string->bytes/utf-8
   (format
    "~a=~a; Path=/; Max-Age=~a; Secure; HttpOnly; SameSite=Strict"
    session-cookie-name
    token
    (auth-manager-session-seconds manager))))

(define (auth-expired-cookie)
  (string->bytes/utf-8
   (format
    "~a=; Path=/; Max-Age=0; Secure; HttpOnly; SameSite=Strict"
    session-cookie-name)))

(define (make-auth-manager user-pairs
                           #:local-networks
                           [local-network-values
                            '("127.0.0.0/8" "::1/128"
                              "10.0.0.0/8" "172.16.0.0/12"
                              "192.168.0.0/16")]
                           #:trusted-proxies
                           [trusted-proxy-values '("127.0.0.0/8" "::1/128")]
                           #:session-seconds [session-seconds 43200])
  (unless (exact-positive-integer? session-seconds)
    (raise-argument-error 'make-auth-manager "exact-positive-integer?"
                          session-seconds))
  (define users (make-hash))
  (for ((entry (in-list user-pairs)))
    (unless (and (pair? entry)
                 (string? (car entry))
                 (string? (cdr entry)))
      (raise-argument-error
       'make-auth-manager
       "(listof (cons/c string? string?))"
       user-pairs))
    (when (string=? (string-trim (car entry)) "")
      (raise-arguments-error
       'make-auth-manager
       "username must not be empty"
       "username" (car entry)))
    (unless (regexp-match? #px"^[$]argon2id[$]" (cdr entry))
      (raise-arguments-error
       'make-auth-manager
       "user password is not an Argon2id hash"
       "username" (car entry)))
    (hash-set! users (string-downcase (string-trim (car entry)))
               (cdr entry)))
  (auth-manager users
                (map parse-network local-network-values)
                (map parse-network trusted-proxy-values)
                session-seconds
                (make-hash)
                (make-hash)
                (make-semaphore 1)))

(module+ test
  (require net/url
           rackunit
           racket/promise
           web-server/http/request-structs)

  (define test-hash (make-password-hash "correct horse battery staple"))
  (check-true (password-hash-valid? "correct horse battery staple" test-hash))
  (check-false (password-hash-valid? "incorrect password" test-hash))

  (define manager
    (make-auth-manager
     (list (cons "Hans" test-hash))
     #:local-networks '("192.168.1.0/24")
     #:trusted-proxies '("127.0.0.1/32")))

  (define (test-request peer [headers '()])
    (request #"GET" (string->url "http://example.test/api/state")
             headers (delay '()) #f "127.0.0.1" 80 peer))

  (check-true (auth-request-local? manager (test-request "192.168.1.42")))
  (check-false (auth-request-local? manager (test-request "192.168.2.42")))
  (check-true
   (auth-request-local?
    manager
    (test-request "127.0.0.1"
                  (list (header #"X-Forwarded-For" #"198.51.100.2, 192.168.1.8")))))
  (check-false
   (auth-request-local?
    manager
    (test-request "198.51.100.2"
                  (list (header #"X-Forwarded-For" #"192.168.1.8")))))

  (define remote (test-request "198.51.100.2"))
  (define token
    (auth-login! manager remote "hans" "correct horse battery staple"))
  (check-true (string? token))
  (define authenticated
    (test-request
     "198.51.100.2"
     (list
      (header #"Cookie"
              (string->bytes/utf-8
               (format "~a=~a" session-cookie-name token))))))
  (check-equal? (auth-request-user manager authenticated) "hans")
  (auth-logout! manager authenticated)
  (check-false (auth-request-user manager authenticated))
  (check-false (auth-login! manager remote "hans" "wrong password")))