;;; Copyright © 2019 - 2022 Jakob L. Kreuze <zerodaysfordays@sdf.org>
;;;
;;; This program is free software; you can redistribute it and/or
;;; modify it under the terms of the GNU General Public License as
;;; published by the Free Software Foundation; either version 3 of the
;;; License, or (at your option) any later version.
;;;
;;; This program is distributed in the hope that it will be useful,
;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
;;; General Public License for more details.
;;;
;;; You should have received a copy of the GNU General Public License
;;; along with this program. If not, see
;;; <http://www.gnu.org/licenses/>.

(define-module (jakob dynamic captcha)
  #:use-module (gcrypt base16)
  #:use-module (gcrypt base64)
  #:use-module (gcrypt hash)
  #:use-module (gcrypt random)
  #:use-module (ice-9 binary-ports)
  #:use-module (ice-9 iconv)
  #:use-module (ice-9 iconv)
  #:use-module (ice-9 local-eval)
  #:use-module (ice-9 match)
  #:use-module (ice-9 popen)
  #:use-module (ice-9 threads)
  #:use-module (jakob dynamic errors)
  #:use-module (json)
  #:use-module (rnrs bytevectors)
  #:use-module ((rnrs base) #:select (assert))
  #:use-module (rnrs conditions)
  #:use-module (rnrs exceptions)
  #:use-module (srfi-197)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-9)
  #:use-module (srfi srfi-11)
  #:use-module (srfi srfi-19)
  #:use-module (srfi srfi-35)
  #:export (make-queue
            id-queue-free
            id-queue-allocated
            release-id!
            dequeue-id!

            new-captcha!
            validate-captcha!
            make-captcha-challenge!

            validate-proof-of-work!
            make-pow-challenge!))

(define-record-type <id-queue>
  (make-id-queue mutex min-free-threshold free-ids allocated-ids)
  id-queue?
  (mutex              id-queue-mutex)
  (min-free-threshold id-queue-min-free-threshold)
  (free-ids           id-queue-free      set-id-queue-free!)
  (allocated-ids      id-queue-allocated set-id-queue-allocated!))

(define* (make-queue n #:key (min-free-threshold 32))
  "Construct a stateful queue for tracking captcha IDs

The parameter N specifies how many free IDs should initially be allocated. The
optional keyword argument MIN-FREE-THRESHOLD specifies when `dequeue-id!' should
iterate through the allocated list and free anything exceeding an
internally-defined `time-to-live-seconds'."
  (make-id-queue (make-mutex) min-free-threshold (iota n) (list)))

(define (append-to-free-queue! id queue)
  "Add ID to the end of the free list of QUEUE"
  (set-id-queue-free!
   queue
   (append! (id-queue-free queue) (list id))))

(define (remove-from-free-queue! id queue)
  "Remove ID from the free list of QUEUE"
  (set-id-queue-free! queue (delete! id (id-queue-free queue))))

(define (append-to-allocated-queue! id queue)
  "Add ID to the end of the allocated list of QUEUE"
  (set-id-queue-allocated!
   queue
   (append! (id-queue-allocated queue) (list (list id (current-time))))))

(define (remove-from-allocated-queue! id queue)
  "Remove ID from the allocated list of QUEUE"
  (set-id-queue-allocated!
   queue
   (filter! (lambda (x) (not (equal? id (car x))))
            (id-queue-allocated queue))))

(define (release-id! id queue)
  "Release ID to the free list of QUEUE"
  (with-mutex (id-queue-mutex queue)
    (assert (find (lambda (x) (equal? id (car x))) (id-queue-allocated queue)))
    (assert (not (member id (id-queue-free queue))))
    (append-to-free-queue! id queue)
    (remove-from-allocated-queue! id queue)))

(define (dequeue-id! queue)
  "Draw a random ID from QUEUE and mark it as allocated"
  (define time-to-live-seconds (* 20 60))
  (with-mutex (id-queue-mutex queue)
    ;; Initial pass to "unintrusively" free any stale IDs.
    (when (< (length (id-queue-free queue))
             (id-queue-min-free-threshold queue))
      (for-each
       (match-lambda
         ((id created-time)
          (when (>= (- (time-second (current-time))
                       (time-second created-time))
                    time-to-live-seconds)
            (remove-from-allocated-queue! id queue))))
       (list-copy (id-queue-allocated queue))))
    ;; If we're still over the threshold, we'll need to be more intrusive.
    ;; Ideally, this is avoided by rate-limiting.
    (when (< (length (id-queue-free queue))
             (id-queue-min-free-threshold queue))
      (let* ((to-take (- (id-queue-min-free-threshold queue)
                         (length (id-queue-free queue))))
             (to-free (map car (take (id-queue-allocated queue) to-take))))
        (set-id-queue-allocated! queue (drop (id-queue-allocated queue) to-take))
        (set-id-queue-free! queue (append! (id-queue-free queue) to-free))))
    (let* ((n (random (length (id-queue-free queue))))
           (id (list-ref (id-queue-free queue) n)))
      (remove-from-free-queue! id queue)
      (append-to-allocated-queue! id queue)
      id)))



(define proc-mutex (make-mutex))
(define tex-challenge-id-queue (make-queue 1024))
(define tex-challenges (make-hash-table 1024))

(define (random-term)
  (match (random 5)
    (0 `(* ,(+ 1 (random 10)) x))
    (1 `(* ,(+ 1 (random 10)) (expt x ,(random 10))))
    (2 `(* ,(+ 1 (random 10)) (exp x)))
    (3 `(* ,(+ 1 (random 10)) (cos x)))
    (4 `(* ,(+ 1 (random 10)) (sin x)))))

(define (sexp->latex sexp)
  (match sexp
    (('+ rest ...) (string-join (map sexp->latex rest) " + "))
    (('* rest ...) (string-join (map sexp->latex rest) " \\cdot "))
    (('sin term) (format #f "\\sin(~a)" (sexp->latex term)))
    (('cos term) (format #f "\\cos(~a)" (sexp->latex term)))
    (('expt term n) (format #f "~a^{~a}" (sexp->latex term) (sexp->latex n)))
    (('exp term) (format #f "e^{~a}" (sexp->latex term)))
    ('x "x")
    (n (cond ((and (number? n) (positive? n)) (format #f "~a" n))
             ((and (number? n) (negative? n)) (format #f "(~a)" n))
             ((number? n) "0")
             (else (error "Do not know how to convert to latex." n))))))

(define (differentiate-sexp sexp)
  (match sexp
    (('+ rest ...) `(+ ,@(map differentiate-sexp rest)))
    (('* coeff term) (if (number? coeff)
                         `(* ,coeff ,(differentiate-sexp term))
                         (error "Do not know how to differentiate.")))
    (('sin term) `(* ,(differentiate-sexp term) (cos ,term)))
    (('cos term) `(* -1 ,(differentiate-sexp term) (sin ,term)))
    (('exp term) `(* ,(differentiate-sexp term) (exp ,term)))
    (('expt term n) `(* ,n (expt ,term ,(- n 1))))
    ('x 1)
    (n (if (number? n)
           0
           (error "Do not know how to differentiate." n)))))

(define (simplify-sexp sexp)
  (match sexp
    (('+ rest ...) `(+ ,@(map simplify-sexp rest)))
    (('* 1 term) (simplify-sexp term))
    (('* 1 rest ...) (simplify-sexp `(* ,@rest)))
    (('sin term) `(sin ,(simplify-sexp term)))
    (('sin term) `(cos ,(simplify-sexp term)))
    (('exp term) `(exp ,(simplify-sexp term)))
    (('expt term 1) (simplify-sexp term))
    (('expt term n) `(expt ,(simplify-sexp term) ,(simplify-sexp n)))
    (term term)))

(define (random-expression)
  (let ((n-terms (+ 2 (random 3))))
    `(+ ,@(map (lambda (x) (random-term)) (iota n-terms)))))

(define (latex->image src)
  (chdir "/tmp")
  (with-mutex proc-mutex
    (call-with-output-file "formula.tex"
      (lambda (port)
        (format port "\\def\\formula{~a}
\\documentclass[border=2pt]{standalone}
\\usepackage{amsmath}
\\usepackage{varwidth}
\\begin{document}
\\begin{varwidth}{\\linewidth}
\\[ \\formula \\]
\\end{varwidth}
\\end{document}
" src)))
    (unless (eqv? 0 (status:exit-val (system "latex formula.tex")))
      (error "Cannot generate DVI" #f))
    (unless (eqv? 0 (status:exit-val (system "dvipng -D 300 formula.dvi")))
      (error "Cannot generate PNG" #f))
    (call-with-input-file "formula1.png" get-bytevector-all)))

(define (new-captcha!)
  (let* ((lower-bound (random 10))
         (upper-bound (+ lower-bound 1 (random 9)))
         (expression (random-expression))
         (latex-src (sexp->latex (simplify-sexp (differentiate-sexp expression))))
         (solution (- (local-eval expression (let ((x upper-bound)) (the-environment)))
                      (local-eval expression (let ((x lower-bound)) (the-environment)))))
         (id (dequeue-id! tex-challenge-id-queue)))
    (hash-set! tex-challenges id solution)
    (values id
            (latex->image (format #f "\\int_{~a}^{~a} ~a \\, dx"
                                  lower-bound
                                  upper-bound
                                  latex-src)))))

(define (validate-captcha! user-answer id)
  (define epsilon 0.1)
  (let ((solution (hash-ref tex-challenges id))
        (id-allocated (not (member id (id-queue-free tex-challenge-id-queue)))))
    ;; FIXME: The predictable IDs means that its' easy for someone to screw with
    ;; someone elses' captcha challenge (by invalidating it before they can
    ;; submit it). Given the combination of our reaping algorithm and
    ;; rate-limiting, does it make sense to only release the ID when the
    ;; response is correct?
    (when (and solution id-allocated)
      (release-id! id tex-challenge-id-queue))
    (and solution
         id-allocated
         (<= (abs (- solution (string->number user-answer)))
             epsilon))))

(define (make-captcha-challenge! request body)
  "API endpoint handler for requesting a captcha challenge"
  (let-values (((challenge-id image) (new-captcha!)))
    (values '((content-type . (application/json)))
            (scm->json-string
             `((challenge-id . ,challenge-id)
               (image . ,(format #f "data:image/jpeg;charset=utf-8;base64,~a"
                                 (base64-encode image))))))))



(define pow-challenge-id-queue (make-queue 1024))
(define pow-challenges (make-hash-table 1024))

;; How many zeroes the SHA-256 hash has to be prefixed by to be a valid proof of work.
(define %hardness 4)

(define (new-proof-of-work-challenge!)
  (let ((id (dequeue-id! pow-challenge-id-queue))
        (challenge (base64-encode (gen-random-bv 32))))
    (hash-set! pow-challenges id challenge)
    (values id challenge)))

(define (validate-proof-of-work! prefix challenge-id)
  (define zero-prefix (string-join (map (lambda (_) "0") (iota %hardness)) ""))
  (when (member challenge-id (id-queue-free pow-challenge-id-queue))
    (panic "No such challenge ID"))
  (let* ((challenge (hash-ref pow-challenges challenge-id))
         (hash-value (chain (list prefix challenge)
                            (string-concatenate _)
                            (string->bytevector _ "utf8")
                            (bytevector-hash _ (lookup-hash-algorithm 'sha256))
                            (bytevector->base16-string _))))
    ;; Invariant from `unless' form:
    ;; (not (member challenge-id (id-queue-free pow-challenge-id-queue)))
    (when challenge
      (release-id! challenge-id pow-challenge-id-queue))
    (and (= 32 (string-length prefix))
         (string-prefix? zero-prefix hash-value))))

(define (make-pow-challenge! request body)
  "API endpoint handler for requesting a proof-of-work challenge"
  (let-values (((challenge-id nonce) (new-proof-of-work-challenge!)))
    (values '((content-type . (application/json)))
            (scm->json-string
             `((hardness . ,%hardness)
               (challenge-id . ,challenge-id)
               (nonce . ,nonce))))))