;; The premise:
;;
;; Better password storage for my various accounts.
;;
;; What are our requirements:
;;
;; Accounts are stored in a VCS-friendly manner: eg, one file per
;; account.
;;
;; We don't need to hide that an account exists, so the account name
;; can be in cleartext (eg, used as a filename).
;;
;; We want multiple levels of security: some passwords are more secret
;; than others and should use a different key.
;;
;; Each account has a number of fields; apart from the account name,
;; they're all encrypted, but each field may be visible or obscured
;; (the latter meaning it's meant to be inserted via the clipboard and
;; never displayed).
;;
;; Design:
;;
;; Accounts are stored in files, with a meaningful name.
;;
;; Inside the file is encrypted data. When decrypted, it's either the
;; details of the account, or another layer of encryption (that may,
;; in turn, contain either account details or more encryption,
;; arbitrarily nested).
;;
;; Account details are a list of (key,value,props) triples: the key is
;; a symbol name for the field, the value is its value (any atomic
;; type), and the props are a list of either symbol flags or (symbol
;; . value) cons cells.
;;
;; "obscured" is a symbol prop that marks fields not to be displayed
;; by default.
;;
;; The decrypted contents of an encrypted blob are either:
;;
;; (account (KEY VALUE [PROPS...])...)
;;
;; or
;;
;; (nested "TAG")BYTES
;;
;; ...where the TAG is the name of the key to prompt the user for, and
;; the BYTES are encrypted bytes.
;;
;; Encryption is a tweetnacl symmetric box, with the nonce
;; prepended. The key is a tweetnacl hash of the user-supplied key
;; string, truncated to keybytes bytes.
;;
;; CLI:
;;
;; pwdb <file> - decrypt and list fields that aren't obscured
;;
;; pwdb -a <file> - decrypt and list all fields
;;
;; pwdb <file> <key>.. - decrypt and dump out raw field value(s), newline separated
;;
;; pwdb -n <File> [tag...] - read (acount ...) from stdin and encrypt; if extra tags are given, nest encryption with them.
;;
;; pwdb -e <file> - spawn $EDITOR to edit non-obscured fields in the file
;;
;; pwdb -e -a <file> - spawn $EDITOR to all fields in the file

;; Base libraries:

(use srfi-4)
(use srfi-13)
(use srfi-69)
(use posix)

;; Required eggs:

(use tweetnacl)
(use args)
(use stty)
(use srfi-27)
(use mwc) ;; from srfi-27
(use entropy-unix) ;; from srfi-27
(use matchable)

;; Core crypto

;; Passphrase cache is used to avoid re-prompting in the edit case,
;; where we decrypt and then re-encrypt. Of course, I'll curse it if I
;; write a "change the passphrase" feature that decrypts with one and
;; re-encrypts with the other...
(define *passphrase-cache* (make-hash-table))

(define (cache-passphrase tag passphrase)
  (hash-table-set! *passphrase-cache* tag passphrase))

(define (get-passphrase tag)
  ;; See if there's a passphrase in the cache
  (if (hash-table-exists? *passphrase-cache* tag)
      (hash-table-ref *passphrase-cache* tag)
      ;; See if there's a passphrase in the environment
      (let ((passphrase-from-env (assoc (if tag (string-append "PWDB_" (symbol->string tag)) "PWDB")
                                      (get-environment-variables))))
        (if passphrase-from-env
            (cdr passphrase-from-env)
            ;; Otherwise, prompt for it
            (with-output-to-file "/dev/tty"
              (lambda ()
                (if tag
                    (printf "Please enter the passphrase for ~a: " tag)
                    (printf "Please enter the passphrase: "))
                (flush-output)
                (let ((key (with-input-from-file "/dev/tty"
                             (lambda ()
                               (with-stty '(not echo) read-line)))))
                  (printf "\n")
                  key)))))))

(define (encrypt1 plaintext key-string nonce-string)
  (let ((key-hash (hash key-string))
        (nonce-hash (hash nonce-string)))
    (assert (>= (string-length key-hash) symmetric-box-keybytes))
    (assert (>= (string-length nonce-hash) symmetric-box-noncebytes))
    (let* ((key (string->blob (substring/shared key-hash 0 symmetric-box-keybytes)))
           (nonce (substring/shared nonce-hash 0 symmetric-box-noncebytes))
           (nonce-u8vector (blob->u8vector/shared (string->blob nonce))))
      (string-append
       nonce
       ((symmetric-box key) plaintext nonce-u8vector)))))

(define (decrypt1 cyphertext key-string)
  (let ((key-hash (hash key-string)))
    (assert (>= (string-length key-hash) symmetric-box-keybytes))
    (assert (>= (string-length cyphertext) symmetric-box-noncebytes))
    (let* ((key (string->blob (substring/shared key-hash 0 symmetric-box-keybytes)))
           (nonce (substring/shared cyphertext 0 symmetric-box-noncebytes))
           (cyphertext-without-nonce (substring/shared cyphertext symmetric-box-noncebytes))
           (nonce-u8vector (blob->u8vector/shared (string->blob nonce))))
      (let ((plaintext ((symmetric-unbox key) cyphertext-without-nonce nonce-u8vector)))
        (if plaintext
            plaintext
            (error "Wrong passphrase!"))))))

#;(assert (string=?
         (decrypt1
          (encrypt1 "Hello, world!" "Key" "Nonce")
          "Key")
         "Hello, world!"))

(define (encrypt plainsexpr nonce-string key-tags)
  ;; Apply tags in reverse order, then final layer
  (foldr (lambda (tag plaintext)
           (let* ((key (get-passphrase tag))
                  (key2 (get-passphrase tag))
                  (_ (unless (string=? key key2) (error "The two passphrases didn't match!")))
                  (cyphertext
                   (encrypt1 plaintext key nonce-string)))
             #;(printf "DEBUG encrypt ~s ~s ~s <- ~s\n" cyphertext key tag plaintext)
             (if tag
                 (string-append
                  (with-output-to-string
                    (lambda ()
                      (write `(nested ,tag))))
                  cyphertext)
                 cyphertext)))
         (with-output-to-string (lambda () (write plainsexpr)))
         (cons #f key-tags)))

;; Returns two values: plainsexpr and list of key-tags
(define (decrypt cyphertext tag)
  (let* ((key (get-passphrase tag))
         (plaintext
          (decrypt1 cyphertext key)))
    (if plaintext
        (cache-passphrase tag key)
        (error "Incorrect passphrase"))
    #;(printf "DEBUG decrypt ~s ~s ~s -> ~s\n" cyphertext key tag plaintext)
    (receive
     (plainsexpr plainbody)
     (with-input-from-string plaintext
       (lambda ()
         (let* ((sexpr (read))
                (rest (read-string)))
           (values sexpr rest))))

     (cond
      ((not (pair? plainsexpr))
       (error "Invalid plaintext"))
      ((eq? (car plainsexpr) 'nested)
       (receive (inner-plain inner-keytags)
                (decrypt plainbody (cadr plainsexpr))
                (values inner-plain (cons tag inner-keytags))))
      (else (values plainsexpr (list tag)))))))

 #;(begin
 (define input '(account))

 (define cyphertext
   (with-input-from-string "A\nB\nC\n"
     (lambda ()
       (encrypt input "Nonce" '(b a)))))

 (define plainsexpr
   (with-input-from-string "C\nB\nA\n"
     (lambda ()
       (decrypt cyphertext #f))))

 (assert (equal? input plainsexpr)))

;; File access

(define (encrypt-and-write plainsexpr path key-tags)
  ;; Path used as nonce
  (let ((cyphertext (encrypt plainsexpr path key-tags)))
    (with-output-to-file path
      (lambda ()
        (change-file-mode path (bitwise-ior perm/irusr perm/iwusr))
        (write-string cyphertext)))))

(define (read-and-decrypt path)
  (receive (plainsexpr key-tags)
           (decrypt
            (with-input-from-file path read-string)
            #f)
           (values plainsexpr key-tags)))

#;(begin
 (define input '(account))

 (with-input-from-string "A\nB\nC\n"
   (lambda ()
     (encrypt-and-write input "test.pwdb" '(b a))))

 (define plainsexpr
   (with-input-from-string "C\nB\nA\n"
     (lambda ()
       (read-and-decrypt "test.pwdb"))))

 (assert (equal? input plainsexpr)))

;; Validity checker

(define (valid-account? account allow-templates? allow-obscured?)
  (and
   (list? account)
   (eq? (car account) 'account)
   (every
    (lambda (field)
      (and
       (list? field)
       (>= (length field) 2)
       (symbol? (first field))
       (or (string? (second field))
           (match
            (second field)
            ('<obscured> allow-obscured?)
            (,('make-passphrase) allow-templates?)
            (,('make-password len) allow-templates?)
            (else #f)))
       (every
        (lambda (prop)
          (or (symbol? prop)
              (and
               (list? prop)
               (symbol? (first prop)))))
        (cddr field))))
    (cdr account))))

;; Accessors

(define (field-key field)
  (first field))

(define (field-value field)
  (second field))

(define (field-props field)
  (cddr field))

(define (obscured? field)
  (member 'obscured (field-props field)))

(define (account-fields account)
  (cdr account))

(define (find-value account key)
  (let loop ((fields (account-fields account)))
    (cond
     ;; No more
     ((null? fields) #f)

     ;; Found it
     ((eq? (field-key (car fields)) key)
      (field-value (car fields)))

     ;; Keep trying
     (else
      (loop (cdr fields))))))

;; Spawn external editor, returns sexpr

(define (edit-file path)
  (let* ((editor (cdr
                  (or (assoc "EDITOR" (get-environment-variables))
                      '("EDITOR" . "vi"))))
         (editor-pid (process-run editor (list path))))
    (process-wait editor-pid)))

(define (edit sexpr valid?)
  (let-values (((fd temp-path) (file-mkstemp "/tmp/pwdb.XXXXXX")))
    (dynamic-wind
        (lambda () #f)
        (lambda ()
         (let ((temp-port (open-output-file* fd)))
           (change-file-mode temp-path (bitwise-ior perm/irusr perm/iwusr))
           (pp sexpr temp-port)
           (close-output-port temp-port)

           ;; Spawn $EDITOR
           (let ((new-sexpr
                  (let loop ()
                    (edit-file temp-path)
                    (let ((new-sexpr (with-input-from-file temp-path read)))
                      (if (valid? new-sexpr)
                          new-sexpr
                          (begin
                            (printf "That's not valid, I'm afraid. Press enter to try again, or ^C to abort.")
                            (read-line)
                            (loop)))))))

             new-sexpr)))
        (lambda ()
          (delete-file temp-path)))))

;; Passphrase generator

(define (random-integer entropy max)
  (let* ((b1 ((entropy-source-u8 entropy)))
         (b2 ((entropy-source-u8 entropy)))
         (u16 (+ (* 256 b1) b2)))
    (inexact->exact (floor (/ (* max u16) 65535)))))

(define (make-passphrase)
  (let* ((gismu (include "gismu.scm"))
         (gismus (vector-length gismu))
         (entropy (make-entropy-source-urandom-device)))

        (string-append
         (vector-ref gismu (random-integer entropy gismus))
         (vector-ref gismu (random-integer entropy gismus))
         (vector-ref gismu (random-integer entropy gismus))
         (vector-ref gismu (random-integer entropy gismus)))))

(define (make-password len)
  (let ((characters
         (list->vector
          (string->list
           ;; Miss out iloO0:;.,| due to confusabilities in bad fonts, and \"' due to quoting pains
           "abcdefghijkmnpqrstuvwxyzABCDEFGHJKLMNPQRSTUVWXYZ123456789%_/+=-*&^%$?!@<>(){}[]~#")))
        (entropy (make-entropy-source-urandom-device)))
    (string-tabulate
     (lambda (_)
       (vector-ref characters (random-integer entropy (vector-length characters))))
     len)))

(define (expand-template tmpl)
  (match
   tmpl
   ;; Check for password generation metacommands
   (,('make-passphrase) (make-passphrase))
   (,('make-password len) (make-password len))
   ;; Otherwise, just recurse
   (else (if (list? tmpl)
             (map expand-template tmpl)
             tmpl))))

;; Interactive mode

(define (copy-to-clipboard clip-command text)
  (receive (child-out child-in child-pid) (process clip-command)
           (write-string text #f child-in)
           (close-output-port child-in)
           (close-input-port child-out)))

(define (interactive-repl account clip-command)
  (printf "index, or quit> ")
  (let* ((response (read-line))
         (response-number (string->number response)))
    (cond
     ;; It's #f if not a valid number
     ((and response-number
           (>= response-number 0)
           (< response-number (length (account-fields account))))
      (let ((field (car (drop (account-fields account) response-number))))
        (copy-to-clipboard clip-command (field-value field))
        (printf "~a copied to clipboard\n" (field-key field)))
      ;; Try again
      (interactive-repl account clip-command))

     ((or
       (string-ci=? response "quit")
       (string-ci=? response "q"))
      ;; Clear clipboard and exit
      (copy-to-clipboard clip-command "")
      (void))

     (else
      (interactive-repl account clip-command)))))

;; Command line engine

(define opts
  (list
   (args:make-option (a all) #:none "Include obscured fields")
   (args:make-option (n new) #:none "Create a new account (with extra layers of encryption named after TAGS)")
   (args:make-option (e edit) #:none "Edit the account (or, with -n, edit a template to create afresh)")
   (args:make-option (i interactive) (required: "COPY-COMMAND") "View the account interactively, with COPY-COMMAND used to copy selected fields to the clipboard")
   (args:make-option (h help) #:none "Display this text" (usage))))

(define (usage)
 (with-output-to-port (current-error-port)
   (lambda ()
     (print "Usage: " (car (argv)) " [options...] FILE [KEY|TAGS...|TEMPLATE-FILE]")
     (newline)
     (print (args:usage opts))))
 (exit 1))

(receive (options operands)
         (args:parse (command-line-arguments) opts)

         (when (zero? (length operands))
                 (usage))

         (let* ((all-mode (assq 'all options))
                (new-mode (assq 'new options))
                (edit-mode (assq 'edit options))
                (view-mode (not (or new-mode edit-mode)))
                (interactive-mode (assq 'interactive options))
                (filename (car operands))
                (rest (cdr operands)))
           (cond
            ;; Check for valid combinations of flags
            ((and new-mode all-mode)
             (usage))
            
            ((and interactive-mode (not view-mode))
             (usage))

            ;; NEW FROM STDIN MODE
            ((and new-mode (not edit-mode))
             (let ((account (read)))
               (unless
                (valid-account? account #t #f)
                (fprintf (current-error-port)
                         "Accounts must be of the form (account (KEY VALUE [PROPS...])...)\n")
                (exit 1))

               (encrypt-and-write
                account
                filename
                (map string->symbol rest))))

            ;; NEW WITH EDIT MODE

            ;; Reading a template file has three functions:

            ;; (1) Providing a template to start from, for convenience.

            ;; (2) Driving password generation parameters.

            ;; (3) Pre-seeding the passphrase cache from decrypting
            ;; the template, so the same keys are used in the new
            ;; file.
            ((and new-mode edit-mode)
             (unless (= 1 (length rest)) (usage))
             (receive (template template-tags)
                      (read-and-decrypt (first rest))
              (let* ((expanded-template (expand-template template))
                     (account (edit expanded-template (cut valid-account? <> #f #f))))

                (encrypt-and-write
                 account
                 filename
                 ;; First element is always #f
                 (cdr template-tags)))))

            ;; VIEW MODE
            (view-mode
             (let ((account (read-and-decrypt filename)))
               (unless (valid-account? account #t #f)
                       (if all-mode
                           (begin
                             (printf "ERROR: Invalid account data:\n")
                             (pp account))
                           (begin
                             (printf "ERROR: Invalid account data (use all-mode to see it)\n")))
                       (exit 1))
               (let ((fields (account-fields account))
                     (desired-keys (map string->symbol rest)))
                 (if (null? desired-keys)
                     ;; List all
                     (begin
                       (let ((idx 0))
                        (for-each
                         (lambda (field)
                           (let ((key (field-key field))
                                 (value (field-value field))
                                 (props (field-props field)))
                             (if (or all-mode
                                     (not (obscured? field)))
                                 (printf "~a: ~a=~s ~s\n" idx key value props)
                                 (printf "~a: ~a=<obscured> ~s\n" idx key props)))
                           (set! idx (+ idx 1)))
                         fields))
                       (when interactive-mode
                             (interactive-repl account (cdr interactive-mode))))
                     ;; List specific keys
                     (for-each
                      (lambda (desired-key)
                        (for-each
                         (lambda (field)
                           (let ((key (field-key field))
                                 (value (field-value field))
                                 (props (field-props field)))
                             (when (eq? key desired-key)
                              (if (or all-mode
                                      (not (obscured? field)))
                                  (printf "~a" value)
                                  (printf "<obscured>"))
                              ;; Pop out a newline if we're producing more than one.
                              ;; For just one, it's nice to have it without a newline, for better piping into other things.
                              (unless (eq? 1 (length desired-keys))
                                      (printf "\n")))))
                         fields))
                      desired-keys)
                     ))))

            ;; EDIT MODE
            (edit-mode
             (let-values (((account key-tags) (read-and-decrypt filename)))
               (if all-mode
                   ;; All mode: Give them everything
                   (let ((new-account (edit account (cut valid-account? <> #t #f))))
                     ;; First element of key-tags is always #f for the outer encryption, so strip that
                     (encrypt-and-write new-account filename (cdr key-tags)))
                   ;; Not all mode: Filter obscured fields, and put them back in afterwards
                   (let* ((new-account (edit
                                       (cons 'account
                                        (map (lambda (field)
                                               (if (obscured? field)
                                                   (append
                                                    (list
                                                     (field-key field)
                                                     '<obscured>)
                                                    (field-props field))
                                                   field))
                                             (account-fields account))) valid-account?))
                          ;; Merge obscured fields back in from the original, unless deleted/overwritten
                          (merged-account
                           (cons 'account
                                 (map (lambda (field)
                                        (if (eq? '<obscured> (field-value field))
                                            (append
                                             (list
                                              (field-key field)
                                              (or (find-value account (field-key field)) "") ; Get previous value
                                              (field-props field)))
                                            field))
                                      (account-fields new-account)))))
                     ;; First element of key-tags is always #f for the outer encryption, so strip that
                     (encrypt-and-write new-account filename (cdr key-tags)))))))))