#lang syndicate
;;; SPDX-License-Identifier: LGPL-3.0-or-later
;;; SPDX-FileCopyrightText: Copyright © 2011-2021 Tony Garnock-Jones <tonyg@leastfixedpoint.com>

(require bitsyntax)
(require rackunit)
(require syndicate/drivers/tcp)

(require "crypto.rkt")
(require "ssh-numbers.rkt")
(require "ssh-message-types.rkt")
(require "ssh-exceptions.rkt")

(provide (struct-out inbound-packet)
	 (struct-out inbound-credit)
	 (struct-out outbound-packet)
	 (struct-out outbound-byte-credit)
	 (struct-out new-keys)

	 default-packet-limit
	 local-algorithm-list

	 ssh-reader
	 ssh-writer)

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Data definitions
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;; A DecodedPacket is one of the packet structures defined in
;; ssh-message-types.rkt.

;; An InboundPacket is an (inbound-packet Number Bytes
;; Maybe<DecodedPacket> Number) representing a packet read from the
;; socket, its sequence number, and the total number of bytes involved
;; in its reception.
(struct inbound-packet (sequence-number payload message transfer-size) #:prefab)

(struct inbound-credit (amount) #:prefab)

(struct outbound-packet (message) #:prefab)

(struct outbound-byte-credit (amount) #:prefab)

(struct new-keys (is-server?
		  derive-key
		  c2s-enc s2c-enc
		  c2s-mac s2c-mac
		  c2s-zip s2c-zip)
	#:prefab)

(struct crypto-configuration (cipher
			      cipher-description
			      hmac
			      hmac-description)
	#:transparent)

;; Description of a supported cipher.
(struct supported-cipher (name factory key-length block-size iv-length)
	#:transparent)

;; Description of a supported hmac algorithm.
(struct supported-hmac (name factory digest-length key-length)
	#:transparent)

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Parameters
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define default-packet-limit (make-parameter 65536))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Encryption, MAC, and Compression algorithm descriptions and parameters
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;; "none" cipher description.
(define null-cipher-description
  (supported-cipher 'none
		    (lambda (enc? key iv)
		      (lambda (block)
			block))
		    0
		    8 ;; pseudo-block-size for packet I/O
		    0))

;; "none" HMAC function.
(define (null-hmac blob)
  #"")

;; "none" HMAC description.
(define null-hmac-description
  (supported-hmac 'none
		  (lambda (key)
		    (error 'null-hmac-description
			   "Cannot construct null hmac instance"))
		  0
		  0))

(define (make-cipher-entry name cipher-spec key-length
                           #:block-size [block-size (cipher-block-size cipher-spec)])
  (list name
	(supported-cipher name
			  (lambda (enc? key iv)
                            (let ((ctx ((if enc? make-encrypt-ctx make-decrypt-ctx)
                                        cipher-spec key iv #:pad #f)))
                              (lambda (input)
                                (cipher-update ctx input))))
                          key-length
                          block-size
			  (cipher-iv-size cipher-spec))))

(define supported-crypto-algorithms
  (list
   (make-cipher-entry 'aes128-ctr '(aes ctr) 16 #:block-size 16)
   (make-cipher-entry 'aes192-ctr '(aes ctr) 24 #:block-size 16)
   (make-cipher-entry 'aes256-ctr '(aes ctr) 32 #:block-size 16)
   (make-cipher-entry 'aes128-cbc '(aes cbc) 16)
   (make-cipher-entry 'aes192-cbc '(aes cbc) 24)
   (make-cipher-entry 'aes256-cbc '(aes cbc) 32)
   (make-cipher-entry '3des-cbc '(des-ede3 cbc) 24)
   )) ;; TODO: actually test these!

(define (make-hmac-entry name digest-spec key-length-or-false)
  (let* ((digest-length (digest-size digest-spec))
	 (key-length (or key-length-or-false digest-length)))
    (list name
	  (supported-hmac name
			  (lambda (key)
			    (lambda (blob)
			      (hmac digest-spec key blob)))
			  digest-length
			  key-length))))

(define supported-hmac-algorithms
  (list (make-hmac-entry 'hmac-sha1 'sha1 #f)))

(define supported-compression-algorithms '(none)) ;; TODO: zlib, and zlib delayed

(define local-algorithm-list
  (let ((crypto-names (map car supported-crypto-algorithms))
	(mac-names (map car supported-hmac-algorithms)))
    (make-parameter
     (lambda ()
       (ssh-msg-kexinit (crypto-random-bytes 16)
			'(diffie-hellman-group14-sha256)
                        '(ssh-ed25519)
			crypto-names
			crypto-names
			mac-names
			mac-names
			supported-compression-algorithms
			supported-compression-algorithms
			'()
			'()
			#f
			0)))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Cryptographic stream configuration
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define initial-crypto-configuration
  (crypto-configuration #f
			null-cipher-description
			null-hmac
			null-hmac-description))

(define (apply-negotiated-options conn-ds nk is-outbound?)
  (match-define (new-keys is-server?
			  (embedded derive-key)
			  c2s-enc s2c-enc
			  c2s-mac s2c-mac
			  c2s-zip s2c-zip) nk)
  ;; TODO: zip
  ;; TODO: make this less ugly. Compute all the keys, select just the ones we need afterward?
  (define c2s
    ;; c2s true iff stream is serverward
    (if is-server? (not is-outbound?) is-outbound?))
  (define enc (if c2s c2s-enc s2c-enc))
  (define mac (if c2s c2s-mac s2c-mac))
  (define zip (if c2s c2s-zip s2c-zip))

  (define cipher-description
    (cond
     ((assq enc supported-crypto-algorithms) => cadr)
     (else (disconnect-with-error conn-ds
                                  SSH_DISCONNECT_KEY_EXCHANGE_FAILED
				  "Could not find driver for encryption algorithm ~v"
				  enc))))
  (define cipher
    ((supported-cipher-factory cipher-description)
     is-outbound?
     (derive-key (if c2s #"C" #"D") (supported-cipher-key-length cipher-description))
     (derive-key (if c2s #"A" #"B") (supported-cipher-iv-length cipher-description))))

  (define hmac-description
    (cond
     ((assq mac supported-hmac-algorithms) => cadr)
     (else (disconnect-with-error conn-ds
                                  SSH_DISCONNECT_KEY_EXCHANGE_FAILED
				  "Could not find driver for HMAC algorithm ~v"
				  mac))))
  (define hmac
    ((supported-hmac-factory hmac-description)
     (derive-key (if c2s #"E" #"F") (supported-hmac-key-length hmac-description))))

  (crypto-configuration cipher cipher-description
			hmac hmac-description))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Transport utilities
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;; MacFunction Natural Bytes -> Bytes
;; Computes the HMAC trailer for a given blob at the given sequence number.
(define (apply-hmac mac sequence-number packet)
  (mac (bit-string->bytes (bit-string-append (integer->bit-string sequence-number 32 #t)
					     packet))))

(define (check-packet-length! conn-ds actual-length limit block-size)
  (when (> actual-length limit)
    (log-warning (format "Packet of length ~v exceeded our limit of ~v"
			 actual-length
			 limit)))
  (when (> actual-length (* 2 limit))
    ;; TODO: For some reason, OpenSSH seems to occasionally slightly
    ;; exceed the packet size limit! (For example, sending a packet of
    ;; length 65564 when I'm expecting a max of 65536.) So we actually
    ;; enforce twice our actual limit.
    (disconnect-with-error conn-ds
                           0 ;; TODO: better reason code?
			   "Packet of length ~v is longer than packet limit ~v"
			   actual-length
			   limit))
  (when (not (zero? (modulo (+ actual-length 4) block-size)))
    ;; the +4 is because the length sent on the wire doesn't include
    ;; the length-of-length, but the requirements for transmitted
    ;; chunks of data are that they be block-size multiples
    ;; *including* the length-of-length
    (disconnect-with-error conn-ds
                           SSH_DISCONNECT_PROTOCOL_ERROR
			   "Packet of length ~v is not a multiple of block size ~v"
			   actual-length
			   block-size)))

;; Integer PositiveInteger -> Integer
;; Rounds "what" up to the nearest multiple of "to".
(define (round-up what to)
  (* to (quotient (+ what (- to 1)) to)))

(check-equal? (round-up 0 8) 0)
(check-equal? (round-up 1 8) 8)
(check-equal? (round-up 7 8) 8)
(check-equal? (round-up 8 8) 8)
(check-equal? (round-up 9 8) 16)

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Encrypted Packet Input
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define (ssh-reader conn-ds conn transfer-control)
  (define input-handler #f)
  (define (update-input-handler #:on-data proc) (set! input-handler proc))
  (assert-control conn
                  #:on-eof (lambda () (stop-current-facet))
                  #:on-data (lambda (data mode) (input-handler data mode)))
  (send! transfer-control 'transfer-control)

  (define packet-size-limit (default-packet-limit))
  (define sequence-number 0)
  (define remaining-credit 0)

  (define config initial-crypto-configuration)

  (define (current-cipher) (crypto-configuration-cipher config))
  (define (block-size)
    (supported-cipher-block-size (crypto-configuration-cipher-description config)))
  (define (decrypt-chunk chunk) ((or (current-cipher) values) chunk))
  (define (subsequent-block-size) (if (current-cipher) (block-size) 1))
  (define (hmac) (crypto-configuration-hmac config))

  (define (issue-credit)
    (when (positive? remaining-credit)
      (send-packet-credit conn (block-size))))

  (define (handle-packet-header encrypted-packet _mode)
    (define first-block (decrypt-chunk encrypted-packet))
    (define packet-length (integer-bytes->integer first-block #f #t 0 4))
    (check-packet-length! conn-ds packet-length packet-size-limit (subsequent-block-size))
    (define amount-of-packet-in-first-block (- (bytes-length first-block) 4)) ;; not incl length
    (define remaining-to-read (- packet-length amount-of-packet-in-first-block))
    (if (positive? remaining-to-read)
        (begin
          (send-packet-credit conn remaining-to-read)
          (update-input-handler
           #:on-data (lambda (encrypted-packet _mode)
                       (define subsequent-chunk (decrypt-chunk encrypted-packet))
                       (check-hmac (bytes-append first-block subsequent-chunk) packet-length))))
        (check-hmac first-block packet-length)))

  (define (check-hmac packet packet-length)
    (define payload-length (let ((padding-length (bytes-ref packet 4)))
                             (- packet-length padding-length 1)))
    (define computed-hmac-bytes (apply-hmac (hmac) sequence-number packet))
    (define mac-byte-count (bytes-length computed-hmac-bytes))
    (if (positive? mac-byte-count)
        (begin
          (send-packet-credit conn mac-byte-count)
          (update-input-handler
           #:on-data (lambda (received-hmac-bytes _mode)
                       (if (equal? computed-hmac-bytes received-hmac-bytes)
                           (finish-packet mac-byte-count packet-length payload-length packet)
                           (disconnect-with-error/local-info conn-ds
                                                             `((expected-hmac ,computed-hmac-bytes)
                                                               (actual-hmac ,received-hmac-bytes))
                                                             SSH_DISCONNECT_MAC_ERROR
                                                             "Corrupt MAC")))))
        (finish-packet 0 packet-length payload-length packet)))

  (define (finish-packet mac-byte-count packet-length payload-length packet)
    (define bytes-read (+ packet-length mac-byte-count))
    (define payload (subbytes packet 5 (+ 5 payload-length)))
    (update-input-handler #:on-data handle-packet-header)
    (send! conn-ds (inbound-packet sequence-number
                                   payload
                                   (ssh-message-decode payload)
                                   bytes-read))
    (set! sequence-number (+ sequence-number 1))
    (set! remaining-credit (- remaining-credit 1))
    (issue-credit))

  (update-input-handler #:on-data handle-packet-header)

  (at conn-ds
    (when (message (inbound-credit $amount))
      (set! remaining-credit (+ remaining-credit amount))
      (issue-credit))

    (when (message ($ nk (new-keys _ _ _ _ _ _ _ _)))
      (set! config (apply-negotiated-options conn-ds nk #f)))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Encrypted Packet Output
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define (ssh-writer conn-ds conn)
  (define config initial-crypto-configuration)
  (define sequence-number 0)

  (define (block-size)
    (supported-cipher-block-size (crypto-configuration-cipher-description config)))
  (define (encrypt-chunk chunk) ((or (crypto-configuration-cipher config) values) chunk))
  (define (hmac) (crypto-configuration-hmac config))

  (at conn-ds
    (when (message (outbound-packet $message))
      (define pad-block-size (block-size))
      (define payload (ssh-message-encode message))
      ;; There must be at least 4 bytes of padding, and padding needs to
      ;; make the packet length a multiple of pad-block-size.
      (define unpadded-length (+ 4 ;; length of length
                                 1 ;; length of length-of-padding indicator
                                 (bit-string-byte-count payload)))
      (define min-padded-length (+ unpadded-length 4))
      (define padded-length (round-up min-padded-length pad-block-size))
      (define padding-length (- padded-length unpadded-length))
      (define packet-length (- padded-length 4))
      ;; ^^ the packet length does *not* include itself!
      (define packet (bit-string->bytes
                      (bit-string (packet-length :: integer bits 32)
                                  (padding-length :: integer bits 8)
                                  (payload :: binary)
                                  ((crypto-random-bytes padding-length) :: binary))))
      (define encrypted-packet (encrypt-chunk packet))
      (define computed-hmac-bytes (apply-hmac (hmac) sequence-number packet))
      (define mac-byte-count (bytes-length computed-hmac-bytes))
      (send-data conn encrypted-packet)
      (send-data conn computed-hmac-bytes)
      (send! conn-ds (outbound-byte-credit (+ (bytes-length encrypted-packet)
                                              (bytes-length computed-hmac-bytes))))
      (set! sequence-number (+ sequence-number 1)))

    (when (message ($ nk (new-keys _ _ _ _ _ _ _ _)))
      (set! config (apply-negotiated-options conn-ds nk #t)))))