;;;; wallet.scm: Bitcoin key and transaction management utilities
;;; J. Welsh, October 2017

;;; Knobs
(define (resource path) (string-append "./" path))
;; Public EC precomputations ("make-cache" to generate)
(define cache-file (resource "secp256k1.cache"))
(define rng-file "/dev/urandom")

;;; Libraries
(load (resource "pkg.scm"))

;; Must set correct base-nibbles parameter for your fixnum size (see bn.scm).
;; 64-bit, 3 tag bits -> 7
;; 32-bit, 3 tag bits -> 3
(define bignum (load-pkg (resource "bignum.scm") error 7))
(import bignum
        bn->hex hex->bn bytes->bn bn->fix fix->bn
        bn-zero? bn-even? bn<
        bn+ bn* bn*fix bn-shift bn-divrem
        bn0 bn1)

(define ecdsa (load-pkg (resource "ecdsa.scm") error bignum))
(import ecdsa secp256k1)

(define bit-ops (load-pkg (resource "bit-ops.scm") error))
(import bit-ops hex->bytes bytes->hex)

(define os
  (load-pkg (resource "os.scm") error #f open-subprocess wait-subprocess))

(define hashes (load-pkg (resource "hashes.scm") bit-ops os #f))
(import hashes ripemd160 sha256)

;;; Helpers

(define null-byte (integer->char 0))

(define (left-pad s len char)
  (string-append (make-string (- len (string-length s)) char) s))

(define (pad-hex-256be s) (left-pad s 64 #\0)) ;; 256 bits = 64 nibbles

(define (pad-bytes-256be s) (left-pad s 32 null-byte))

(define (string-head s n) (substring s 0 n))

(define (string-tail s n)
  (let ((len (string-length s)))
    (if (negative? n)
      (substring s (+ n len) len)
      (substring s n len))))

(define (string-reverse s)
  (let ((r (make-string (string-length s)))
        (last (- (string-length s) 1)))
    (do ((i 0 (+ i 1))) ((> i last) r)
      (string-set! r i (string-ref s (- last i))))))

(define (bn->bytes n)
  ;; hacky...
  (hex->bytes (bn->hex n)))

(define (bn-pack-le nbytes)
  (let ((nibbles (* nbytes 2)))
    (lambda (n)
      (string-reverse (hex->bytes (left-pad (bn->hex n) nibbles #\0))))))

(define (string->bytevec s)
  (let* ((len (string-length s))
         (v (make-vector len)))
    (do ((i 0 (+ i 1))) ((= i len) v)
      (vector-set! v i (char->integer (string-ref s i))))))

(define (byte-string->bn s)
  (bytes->bn (string->bytevec s)))

(define (byte-string . bytes)
  (list->string (map integer->char bytes)))

(define (display-line x)
  (display x)
  (newline))

(define (display-lines l)
  (for-each display-line l))

(define (write-to-file f disp-func x)
  (with-output-to-file f (lambda () (disp-func x)))
  (display-line (string-append "wrote " f)))

(define (read-line . port)
  (let loop ((acc '()))
    (let ((c (apply read-char port)))
      (if (char=? c #\Newline) (list->string (reverse acc))
        (if (eof-object? c) (if (null? acc) c
                              (list->string (reverse acc)))
          (loop (cons c acc)))))))

(define (read-priv-key-hex)
  (let ((k (read-line)))
    (if (not (= (string-length k) 64)) (error "ill-formed key"))
    k))

(define (any pred l)
  (and (pair? l)
       (or (pred (car l))
           (any pred (cdr l)))))

;; Wants to be the full-time printer when it grows up.
(define (pprint x)
  (define (pprint x indent)
    (if (and (list? x) (any pair? x))
      (begin
        (write-char #\()
        (pprint-items x (string-append indent (list-indent (car x))))
        (write-char #\)))
      (write x)))
  (define (list-indent first)
    (if (or (pair? first) (null? first))
      " "
      "  "))
  (define (pprint-items l indent)
    (pprint (car l) indent)
    (for-each (lambda (x)
                (newline)
                (display indent)
                (pprint x indent))
              (cdr l)))
  (pprint x "")
  (newline))

;;; Base58 [TRB/base58.h]
;;; (See also http://en.bitcoin.it/wiki/Base58Check_encoding)

(define bn58 (fix->bn 58))

(define (b58digit i) ;; alphanumeric - 0OIl
  (string-ref "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz" i))

(define b58digit->integer
  (let ((inverse (make-vector 256 #f)))
    (do ((i 0 (+ i 1))) ((= i 58))
      (vector-set! inverse (char->integer (b58digit i)) i))
    (lambda (d)
      (or (vector-ref inverse (char->integer d))
          (error "invalid base58 digit:" d)))))

(define (encode-base58 data)
  (let loop ((n (byte-string->bn data))
             (acc '()))
    (if (not (bn-zero? n))
      (bn-divrem n bn58 (lambda (q r)
                          (loop q (cons (bn->fix r) acc))))
      ;; Leading zero bytes pass through one-to-one
      (do ((i 0 (+ i 1))
           (acc acc (cons 0 acc)))
        ((or (= i (string-length data))
             (not (zero? (char->integer (string-ref data i)))))
         (list->string (map b58digit acc)))))))

(define (decode-base58 data)
  (do ((digits (map b58digit->integer (string->list data)) (cdr digits))
       (zeros 0 (+ 1 zeros)))
    ((or (null? digits) (not (zero? (car digits))))
     (do ((acc bn0 (bn+ (bn* acc bn58) (fix->bn (car digits))))
          (digits digits (cdr digits)))
       ((null? digits) (string-append (make-string zeros null-byte)
                                      (if (bn-zero? acc) ""
                                        (bn->bytes acc))))))))

;; Well-known Base58Check payload type tags
(define b58-version-p2pkh 0)
(define b58-version-p2sh 5)
(define b58-version-secret 128)
(define b58-version-testnet-p2pkh 111)
(define b58-version-testnet-p2sh 196)
(define b58-version-testnet-secret 239)

(define (encode-base58check version data)
  (let ((data (string-append (byte-string version) data)))
    (encode-base58
      (string-append data (string-head (sha256 (sha256 data)) 4)))))

(define (decode-base58check data)
  (let* ((data (decode-base58 data))
         (cut (- (string-length data) 4)))
    (if (negative? cut) (error "decode-base58check: checksum too short"))
    (let ((payload (string-head data cut))
          (check (string-tail data cut)))
      (if (not (string=? check (string-head (sha256 (sha256 payload)) 4)))
        (error "decode-base58check: bad checksum"))
      payload)))

;;; OpenSSL binary encodings

;; EC point (public key) [1]
;; Leading byte encodes form:
;; (0)     the point at infinity
;; (2 x)   compressed, even y
;; (3 x)   compressed, odd y
;; (4 x y) uncompressed
;; (6 x y) hybrid, even y
;; (7 x y) hybrid, odd y
(define (encode-point p form)
  (if (eq? p 'inf) (byte-string 0)
    (let ((x (car p))
          (y (cdr p))
          (pack (lambda (n) (pad-bytes-256be (bn->bytes n)))))
      (case form
        ((compressed)
         (string-append (byte-string (if (bn-even? y) 2 3)) (pack x)))
        ((uncompressed)
         (string-append (byte-string 4) (pack x) (pack y)))
        ((hybrid)
         ;; Hybrid form seems entirely pointless to me, but is included for
         ;; completeness.
         (string-append (byte-string (if (bn-even? y) 6 7))
                        (pack x) (pack y)))
        (else (error "encode-point: bad form:" form))))))

;; DERP encoding: a subset of the Distinguished Encoding Rules of Abstract
;; Syntax Notation One

(define (derp-encapsulate tag constructed? contents)
  ;; Class assumed to be 0 (universal)
  ;; Tag assumed to be 0-30
  ;; Length assumed to be 0-127
  (string-append (byte-string (if constructed? (+ tag 32) tag)
                              (string-length contents))
                 contents))

(define (derp-contents s)
  (string-tail s 2))

(define (encode-derp-integer n)
  ;; INTEGER is encoded base 256, two's complement, big endian, with the
  ;; minimal number of octets, except zero which is a single zero octet
  (let* ((b (bn->bytes n))
         (pad? (or (zero? (string-length b))
                   (> (char->integer (string-ref b 0)) 127))))
    (derp-encapsulate 2 #f (if pad? (string-append (byte-string 0) b) b))))

(define (decode-derp-integer s start)
  (byte-string->bn (derp-contents s)))

(define (encode-derp-sequence . items)
  ;; SEQUENCE is encoded as the concatenation of the respective encoded items
  (derp-encapsulate 16 #t (apply string-append items)))

(define (encode-derp-sig r s)
  ;; Per [SSL/crypto/ecdsa/ecs_asn1.c], [SSL/crypto/asn1/x_bignum.c]
  (encode-derp-sequence (encode-derp-integer r) (encode-derp-integer s)))

;; Might be bloat, given that a detached wallet can't validate transactions,
;; but might come in handy
(define (parse-derp s)
  (define (parse-item start) ;; returns pair (bytes parsed, value)
    (let ((ident (char->integer (string-ref s start)))
          (len (char->integer (string-ref s (+ 1 start))))
          (start (+ start 2)))
      (if (= (remainder ident 32) 31)
        (error "unsupported DER high-tag-number form"))
      (if (> len 127) (error "unsupported DER multi-byte length"))
      (let ((end (+ start len)))
        (if (> end (string-length s)) (error "DER length out of bounds"))
        (cons (+ len 2)
              (case ident
                ((2) ;; INTEGER
                 (byte-string->bn (substring s start (+ start len))))
                ((48) ;; SEQUENCE
                 (let loop ((start start) (acc '()))
                   (if (= start end) (reverse acc)
                     (let ((p (parse-item start)))
                       (loop (+ start (car p)) (cons (cdr p) acc))))))
                (else (error "unsupported DER type" ident)))))))
  (let ((p (parse-item 0)))
    (if (not (= (car p) (string-length s)))
      (error "excess bytes after DER parsing"))
    (cdr p)))

;;; Bitcoin address encoding

(define (point->address p)
  (encode-base58check b58-version-p2pkh
                      (ripemd160 (sha256 (encode-point p 'compressed)))))

;;; Bitcoin transactions

(define (alist fields)
  (lambda vals
    (map list fields vals)))

(define (alist-get l field . default)
  (let ((l (assq field l)))
    (if l (cadr l)
      (if (null? default) (error "field not found" field)
        (car default)))))

(define make-input (alist '(txid index value sig-script)))
(define make-output (alist '(value address)))
(define make-transaction (alist '(inputs outputs)))

;;; Script, per [TRB/script.h]

(define (enumerate start items . more)
  (let loop ((i start) (l items) (more more) (acc '()))
    (if (null? l) (if (null? more) (reverse acc)
                    (loop (car more) (cadr more) (cddr more) acc))
      (loop (+ i 1) (cdr l) more (cons (list (car l) i) acc)))))

(define script-ops
  (enumerate
    0
    '("0") ;aka FALSE
    ;1-75 indicate length of following string to push
    76
    '("PUSHDATA1" "PUSHDATA2" "PUSHDATA4" "1NEGATE" "RESERVED"
      "1" "2" "3" "4" "5" "6" "7" "8" "9" "10" "11" "12" "13" "14" "15" "16"
      ;control
      "NOP" "VER" "IF" "NOTIF" "VERIF" "VERNOTIF" "ELSE" "ENDIF" "VERIFY"
      "RETURN"
      ;stack ops
      "TOALTSTACK" "FROMALTSTACK" "2DROP" "2DUP" "3DUP" "2OVER" "2ROT" "2SWAP"
      "IFDUP" "DEPTH" "DROP" "DUP" "NIP" "OVER" "PICK" "ROLL" "ROT" "SWAP"
      "TUCK"
      ;splice ops
      "CAT" "SUBSTR" "LEFT" "RIGHT" "SIZE"
      ;bit logic
      "INVERT" "AND" "OR" "XOR" "EQUAL" "EQUALVERIFY" "RESERVED1" "RESERVED2"
      ;numeric
      "1ADD" "1SUB" "2MUL" "2DIV" "NEGATE" "ABS" "NOT" "0NOTEQUAL" "ADD" "SUB"
      "MUL" "DIV" "MOD" "LSHIFT" "RSHIFT" "BOOLAND" "BOOLOR" "NUMEQUAL"
      "NUMEQUALVERIFY" "NUMNOTEQUAL" "LESSTHAN" "GREATERTHAN"
      "LESSTHANOREQUAL" "GREATERTHANOREQUAL" "MIN" "MAX" "WITHIN"
      ;crypto
      "RIPEMD160" "SHA1" "SHA256" "HASH160" "HASH256" "CODESEPARATOR"
      "CHECKSIG" "CHECKSIGVERIFY" "CHECKMULTISIG" "CHECKMULTISIGVERIFY"
      ;expansion
      "NOP1" "NOP2" "NOP3" "NOP4" "NOP5" "NOP6" "NOP7" "NOP8" "NOP9" "NOP10")
    253 ;template matching params
    '("PUBKEYHASH" "PUBKEY" "INVALIDOPCODE")))

;; Human-readable, bijective script decoding
(define (encode-script s)
  (define (encode-op name)
    (byte-string (cadr (or (assoc name script-ops)
                           (error "bad opcode name" name)))))
  (let loop ((s s) (acc '()))
    (if (null? s) (apply string-append (reverse acc))
      (let ((op (car s)))
        (if (integer? op)
          (loop (cddr s) (cons (hex->bytes (cadr s))
                               (cons (byte-string op) acc)))
          (let ((acc (cons (encode-op op) acc)))
            (if (member op '("PUSHDATA1" "PUSHDATA2" "PUSHDATA4"))
              (let* ((data (hex->bytes (cadr s)))
                     (len (pack-le (cond ((equal? op "PUSHDATA1") 1)
                                         ((equal? op "PUSHDATA2") 2)
                                         ((equal? op "PUSHDATA4") 4))
                                   (string-length data))))
                (loop (cddr s) (cons data (cons len acc))))
              (loop (cdr s) acc))))))))

(define (script-push bytes)
  (let ((len (string-length bytes)))
    (list (cond ((< len 76) len)
                ((< len (expt 2 8)) "PUSHDATA1")
                ((< len (expt 2 16)) "PUSHDATA2")
                ((< len (expt 2 32)) "PUSHDATA4")
                (else (error "script-push overflow")))
          (bytes->hex bytes))))

(define (output-script address)
  (let ((address (decode-base58check address)))
    (if (not (= (string-length address) 21)) (error "bad address length"))
    (let ((version (char->integer (string-ref address 0)))
          (hash (string-tail address 1)))
      (if (or (= version b58-version-p2pkh)
              (= version b58-version-testnet-p2pkh))
        `("DUP" "HASH160" ,@(script-push hash) "EQUALVERIFY" "CHECKSIG")
        (if (or (= version b58-version-p2sh)
                (= version b58-version-testnet-p2sh))
          `("HASH160" ,@(script-push hash) "EQUAL")
          (error "bad address type"))))))

(define sighash-all 1)

(define (input-script r s pubkey form)
  (append (script-push (string-append (encode-derp-sig r s)
                                      (byte-string sighash-all)))
          (script-push (encode-point pubkey form))))

;;; Wire protocol
;;; (See also http://en.bitcoin.it/wiki/Protocol_documentation)

(define (pack-le nbytes n)
  (do ((i 0 (+ i 1))
       (n n (quotient n 256))
       (bytes '() (cons (integer->char (remainder n 256)) bytes)))
    ((= i nbytes)
     (if (zero? n) (list->string (reverse bytes))
       (error "pack-le: overflow" nbytes)))))

;; Encode a bignum under 2^64 in Bitcoin's variable-length integer format
(define encode-var-int
  (let ((pack (lambda (tag nbytes)
                (let ((tag (byte-string tag))
                      (pack-n (bn-pack-le nbytes)))
                  (lambda (n) (string-append tag (pack-n n)))))))
    (let ((c1 (fix->bn 253))
          (c2 (bn-shift bn1 16))
          (c3 (bn-shift bn1 32))
          (c4 (bn-shift bn1 64))
          (pack2 (pack 253 2))
          (pack4 (pack 254 4))
          (pack8 (pack 255 8)))
      (lambda (n)
        (cond ((bn< n c1) (byte-string (bn->fix n)))
              ((bn< n c2) (pack2 n))
              ((bn< n c3) (pack4 n))
              ((bn< n c4) (pack8 n))
              (else (error "encode-var-int overflow")))))))

(define (encode-var-str s)
  (string-append (encode-var-int (fix->bn (string-length s))) s))

;;; Commands

(define (pop-arg)
  (if (null? *args*) (error "too few arguments"))
  (let ((a (car *args*)))
    (set! *args* (cdr *args*))
    a))

(define cache (delay ((secp256k1 'read-cache) cache-file)))

(define (make-cache) ((secp256k1 'write-cache) cache-file))

;; Generate a private key and its addresses, storing to files
(define (gen-key)
  (force cache)
  (if (null? *args*) (error "must specify key file basename"))
  (let* ((basename (pop-arg))
         (priv (call-with-input-file rng-file (secp256k1 'gen-priv-key))))
    (write-to-file (string-append basename ".k") display-line
                   (pad-hex-256be (bn->hex priv)))
    (write-to-file (string-append basename ".a") display-line
                   (point->address ((secp256k1 'priv->pub) priv)))))

;; Print addresses for a private key fed to standard input
(define (priv2addr)
  (force cache)
  (display-line (point->address ((secp256k1 'priv->pub)
                                 (hex->bn (read-priv-key-hex))))))

;; Print addresses for a curve point fed as S-exp to standard input
(define (point2addr)
  (let ((p (read)))
    (display-line (point->address
                    (if (eq? p 'inf) p
                      (cons (hex->bn (car p))
                            (hex->bn (cdr p))))))))

;; Print WIF for a private key fed to standard input, both with and without the
;; compression flag
(define (priv2wif)
  (let ((key-bytes (hex->bytes (read-priv-key-hex))))
    (display-lines
      (list (encode-base58check b58-version-secret
                                (string-append key-bytes (byte-string 1)))
            (encode-base58check b58-version-secret key-bytes)))))

(define (output-to)
  (let* ((addr (read-line))
         (o (make-output 10 addr))
         (s (output-script (alist-get o 'address))))
      (pprint o)
      (pprint s)
      (pprint (bytes->hex (encode-script s)))))

;; Run self-tests
(define (test)
  (define (test-equal a b)
    (if (equal? a b)
      (write 'pass)
      (write `(fail ,a != ,b)))
    (newline))

  ;; From python-bitcoinlib
  (for-each
    (lambda (r)
      (let* ((bytes (hex->bytes (car r)))
             (enc (encode-base58 bytes))
             (dec (decode-base58 enc)))
        (test-equal enc (cadr r))
        (test-equal dec bytes)))
    '(("" "")
      ("61" "2g")
      ("626262" "a3gV")
      ("636363" "aPEr")
      ("73696d706c792061206c6f6e6720737472696e67" "2cFupjhnEsSn59qHXstmK2ffpLv2")
      ("00eb15231dfceb60925886b67d065299925915aeb172c06647" "1NS17iag9jJgTHD1VXjvLCEnZuQ3rJDE9L")
      ("516b6fcd0f" "ABnLTmg")
      ("bf4f89001e670274dd" "3SEo3LWLoPntC")
      ("572e4794" "3EFU7m")
      ("ecac89cad93923c02321" "EJDM8drfXA6uyA")
      ("10c8511e" "Rt5zm")
      ("00000000000000000000" "1111111111"))))

(define (help)
  (display-line "Available commands:")
  (for-each (lambda (c) (display-line (car c))) commands))

;;; CLI dispatch

(define commands
  `((make-cache ,make-cache)
    (gen-key ,gen-key)
    (priv2addr ,priv2addr)
    (point2addr ,point2addr)
    (priv2wif ,priv2wif)
    (output-to ,output-to)
    (test ,test)
    (help ,help)))

(if (not (null? *args*))
  (let ((prog-name (pop-arg)))
    (let ((cmd (if (null? *args*) 'help (string->symbol (pop-arg)))))
      ((cond ((assq cmd commands) => cadr)
             (else (display "Bad command: ") (display-line cmd)
                   help))))))

;;; Notes
;
; [1]
; While apparently simple, point encoding is one of the more dangerous parts
; of the code, as any incompatibility could result in seemingly valid
; addresses that turn out to be unspendable after you've already buried your
; treasure there. (Likewise the hash functions.) Verifying requires quite the
; code walk:
;
; [TRB/key.h] CKey::GetPubKey calls
; [SSL/crypto/ec/ec_asn1.c] i2o_ECPublicKey (which, filename notwithstanding,
;   does not involve ASN.1), which calls
; [SSL/crypto/ec/ec_oct.c] EC_POINT_point2oct, which calls
; [SSL/crypto/ec/ecp_oct.c] ec_GFp_simple_point2oct, in which is found the
;   encoding used here, assuming (group->meth->flags & EC_FLAGS_DEFAULT_OCT)
;   and (group->meth->field_type == NID_X9_62_prime_field). "group" is the
;   group field of the EC_KEY passed to i2o_ECPublicKey, initialized in
; [TRB/key.h] CKey::CKey by calling
; [SSL/crypto/ec/ec_key.c] EC_KEY_new_by_curve_name(NID_secp256k1), which
;   initializes group by calling
; [SSL/crypto/ec/ec_curve.c] EC_GROUP_new_by_curve_name on the same NID,
;   which calls
; [id.] ec_group_new_from_data on the member of the constant curve_list with
;   matching NID. Since the curve_list item in question has a data field of
;   &_EC_SECG_PRIME_256K1.h (an EC_CURVE_DATA constant having field_type of
;   NID_X9_62_prime_field) and meth field of 0, the branch is taken that
;   calls
; [SSL/crypto/ec/ec_cvt.c] EC_GROUP_new_curve_GFp, which initializes meth to
;   either
; [SSL/crypto/ec/ecp_nist.c] EC_GFp_nist_method() or
; [SSL/crypto/ec/ecp_mont.c] EC_GFp_mont_method(), by calling
; [SSL/crypto/ec/ec_lib.c] EC_GROUP_new. In both cases the result is a
;   constant EC_METHOD having the above required flags and field_type,
;   assuming OPENSSL_FIPS is not defined.
;
; Even this does not constitute a proof, due to the possibility of mutation by
; seemingly-unrelated functions.
;
; An implication of this encoding is that even though [TRB] only emits
; uncompressed form, and no implementation has ever to my knowledge used
; hybrid form, there are three P2PKH addresses spendable by any given private
; key, and a consensus-compatible verifier must support all forms.

;;; References
;
; [TRB] The Real Bitcoin, version 0.5.4-RELEASE. http://thebitcoin.foundation/
;
; [SSL] OpenSSL, version 1.0.1g.