#lang racket/base

(require (planet tonyg/bitsyntax))
(require (planet vyzo/crypto:2:3))
(require (planet vyzo/crypto/util)) ;; hex, unhex
(require racket/port)
(require racket/bool)

(require rackunit)
(require "aes-ctr.rkt")
(require "safe-io.rkt")
(require "ssh-numbers.rkt")
(require "ssh-message-types.rkt")
(require "oakley-groups.rkt")

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

;; An exn:fail:contract:protocol, when thrown by the transport (TODO:
;; clarify scope of this behaviour) will cause a SSH_MSG_DISCONNECT to
;; be sent to the remote party with the included reason code, using
;; the exn-message as the description. The local-info field is useful
;; information for diagnosing problems known to the local stack that
;; should not be transmitted to the remote party. For example, upon
;; detection of a MAC failure, it might be useful to know the expected
;; and actual MACs for debugging, but they should not be sent over the
;; wire because we could be experiencing some kind of attack.
(struct exn:fail:contract:protocol exn:fail:contract
	(reason-code local-info originated-at-peer?)
	#:transparent)

;; A RekeyState is one of
;; - a (rekey-wait Number Number), representing a time or
;;   transfer-amount by which rekeying should be started
;; - a (rekey-local SshMsgKexinit), when we've sent our local
;;   algorithm list and are waiting for the other party to send theirs
;; - a (rekey-in-progress KeyExchangeState), when both our local
;;   algorithm list has been sent and the remote one has arrived and the
;;   actual key exchange has begun
(struct rekey-wait (deadline threshold-bytes) #:transparent)
(struct rekey-local (local-algorithms) #:transparent)
(struct rekey-in-progress (state) #:transparent)

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

;; A PacketDispatcher is a Hashtable mapping Byte to PacketHandler.

;; A PacketHandler is a (Bytes DecodedPacket ConnectionState -> ConnectionState).
;; TODO: fix this definition
;; The raw received bytes of the packet are given because sometimes
;; cryptographic operations on the received bytes are mandated by the
;; protocol.

;; A StreamState is a (stream-state Port Encryptor SupportedCipher
;; Uint32 Natural MacFunction SupportedHmac Natural) representing the
;; negotiated and computed state of the packet-delimiting,
;; -encrypting, and -MACing layer. There's one for each direction
;; (inbound and outbound) of a connection.
(struct stream-state (port
		      cipher
		      cipher-description
		      sequence-number ;; TODO: clip to Uint32
		      bytes-transferred
		      hmac
		      hmac-description
		      packet-size-limit)
	#:transparent)

;; A ConnectionState is a (connection StreamState StreamState
;; PacketDispatcher ... TODO fix this) representing the complete state
;; of the SSH transport, authentication, and connection layers.
(struct connection (in
		    out
		    dispatch-table
		    global-request-dispatch-table
		    channel-open-handler
		    rekey-state
		    is-server?
		    local-id
		    remote-id
		    session-id) ;; starts off #f until initial keying
	#:transparent)

;; Generic inputs into the exchange-hash part of key
;; exchange. Diffie-Hellman uses these fields along with the host key,
;; the exchange values, and the shared secret to get the final hash.
(struct exchange-hash-info (client-id
			    server-id
			    client-kexinit-bytes
			    server-kexinit-bytes)
	#: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 identification-recogniser #rx"^SSH-")
(define (identification-line? str)
  (regexp-match identification-recogniser str))

(define required-peer-identification-regex (make-parameter #rx"^SSH-2\\.0-.*"))

(define client-preamble-lines (make-parameter '()))
(define client-identification-string (make-parameter "SSH-2.0-RacketSSH_0.0"))

(define default-packet-limit (make-parameter 65536))
(define rekey-interval (make-parameter 5)) ;;3600))
(define rekey-volume (make-parameter 1000000000))
(define inter-packet-timeout (make-parameter 1)) ;;300))
(define intra-packet-timeout (make-parameter 1)) ;;300))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; 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-evp-cipher-entry name cipher)
  (list name
	(supported-cipher name
			  (lambda (enc? key iv)
			    (let ((state ((if enc? cipher-encrypt cipher-decrypt)
					  cipher key iv #:padding #f)))
			      (lambda (block)
				(cipher-update! state block))))
			  (cipher-key-length cipher)
			  (cipher-block-size cipher)
			  (cipher-iv-length cipher))))

(define (aes-ctr-cipher-factory enc? key iv)
  (let ((state (start-aes-ctr key iv)))
    (lambda (block)
      (aes-ctr-process! state block))))

(define (make-aes-ctr-entry name key-length)
  (list name
	(supported-cipher name
			  aes-ctr-cipher-factory
			  key-length
			  16
			  16)))

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

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

(define supported-hmac-algorithms
  (list (make-hmac-entry 'hmac-md5 digest:md5 #f)
	(make-hmac-entry 'hmac-sha1 digest: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 (random-bytes 16)
			'(diffie-hellman-group14-sha1
			  diffie-hellman-group1-sha1)
			'(ssh-rsa ssh-dss)
			crypto-names
			crypto-names
			mac-names
			mac-names
			supported-compression-algorithms
			supported-compression-algorithms
			'()
			'()
			#f
			0)))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Error signalling
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define (disconnect-with-error reason-code format-string . args)
  (apply disconnect-with-error/local-info '() reason-code format-string args))

(define (disconnect-with-error/local-info local-info reason-code format-string . args)
  (let ((message (apply format format-string args)))
    (raise (exn:fail:contract:protocol message
				       (current-continuation-marks)
				       reason-code
				       local-info
				       #f))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; I/O Utilities for timeouts and decryption
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define (read-bytes/timeout count in timeout)
  (sync/timeout timeout (read-bytes-evt count in)))

(define (read-bytes/decrypt count in timeout decryptor)
  (let ((encrypted (read-bytes/timeout count in timeout)))
    (cond
     ((false? encrypted) #f)
     ((eof-object? encrypted) eof)
     (else (decryptor encrypted)))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Encrypted Packet I/O
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define (check-packet-length! actual-length limit block-size)
  (when (> actual-length limit)
    (disconnect-with-error 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 SSH_DISCONNECT_PROTOCOL_ERROR
			   "Packet of length ~v is not a multiple of block size ~v"
			   actual-length
			   block-size)))

;; TODO: The OpenSSH sshd won't accept a rekeying until authentication
;; is complete, so until we implement the auth layer, we'll get an
;; SSH_MSG_UNIMPLEMENTED when we send SSH_MSG_KEXINIT after the
;; initial keying.

;; TODO: Remove the incredibly short timeouts above (both inter- and
;; intra-packet-timeout, and rekey-interval).

;; StreamState Natural -> StreamState
(define (bump-sequence-number state byte-count)
  (struct-copy stream-state state
	       [sequence-number
		;; It's an unsigned, 32-bit packet counter, so clip it at 32 bits.
		(bitwise-and #xffffffff (+ 1 (stream-state-sequence-number state)))]
	       [bytes-transferred
		(+ byte-count (stream-state-bytes-transferred state))]))

;; ConnectionState Boolean ->
;;    (values EndOfFile EndOfFile ConnectionState)
;; or (values #f #f ConnectionState)
;; or (values Bytes DecodedPacket ConnectionState)
;;
;; Read and decode a transport message from in-state. If it can't be
;; decoded (we don't support that message type), complain with a
;; SSH_MSG_UNIMPLEMENTED packet. Finally, return a quadruple of the
;; packet, the decoded message, the updated input state, and the
;; updated output state. May return eof or #f instead of a packet for
;; end-of-file or timeout, respectively, depending on
;; error-on-eof-or-timeout.
(define (read-message conn [error-on-eof-or-timeout #t])
  (let-values (((packet conn) (read-packet conn error-on-eof-or-timeout)))
    (if (not (bytes? packet))
	(values packet packet conn)
	(let ((message (ssh-message-decode packet)))
	  (write `(received ,message)) (newline) (flush-output)
	  (if message
	      (values packet message conn)
	      (let ((bad-seq-num (most-recent-received-sequence-number conn)))
		;; TODO: remove this debug output
		(display "BAD PACKET ")
		(display (hex packet))
		(newline)
		(flush-output)
		(read-message (write-message! (ssh-msg-unimplemented bad-seq-num) conn)
			      error-on-eof-or-timeout)))))))

;; ConnectionState -> Natural
;; Returns the sequence number of the most recently received packet.
(define (most-recent-received-sequence-number conn)
  (- (stream-state-sequence-number (connection-in conn)) 1))

;; Packet format on the wire:
;; uint32    packet_length
;; byte      padding_length
;; byte[n1]  payload; n1 = packet_length - padding_length - 1
;; byte[n2]  random padding; n2 = padding_length
;; byte[m]   mac (Message Authentication Code - MAC); m = mac_length

;; ConnectionState Boolean -> (values Bytes ConnectionState)
;; Read, MAC-check, and decrypt a single packet from in-state.
(define (read-packet conn error-on-eof-or-timeout)
  (define in-state (connection-in conn))
  (define cipher (stream-state-cipher in-state))
  (define block-size (supported-cipher-block-size (stream-state-cipher-description in-state)))
  (define in (stream-state-port in-state))
  (define first-block-size block-size)
  (define subsequent-block-size (if cipher block-size 1))
  (define decryptor (if cipher cipher values))
  (define first-block (read-bytes/decrypt first-block-size in
					  (inter-packet-timeout) decryptor))
  (cond
   ((false? first-block)
    (if error-on-eof-or-timeout
	(disconnect-with-error SSH_DISCONNECT_CONNECTION_LOST
			       "Timeout waiting for a packet")
	(values #f conn)))
   ((eof-object? first-block)
    (if error-on-eof-or-timeout
	(error 'read-packet "End-of-file at the start of a packet")
	(values first-block conn)))
   (else
    (let* ((packet-length (integer-bytes->integer first-block #f #t 0 4)))
      (check-packet-length! packet-length
			    (stream-state-packet-size-limit in-state)
			    subsequent-block-size)
      (let* ((padding-length (bytes-ref first-block 4))
	     (payload-length (- packet-length padding-length 1))
	     (amount-of-packet-in-first-block (- (bytes-length first-block) 4)) ;; not incl length
	     (remaining-to-read (- packet-length amount-of-packet-in-first-block)))
	(define (read-packet-trailer packet)
	  (let ((bytes-read (+ (check-hmac! (apply-hmac (stream-state-hmac in-state)
							(stream-state-sequence-number in-state)
							packet)
					    in)
			       packet-length)))
	    (values (subbytes packet 5 (+ 5 payload-length))
		    (struct-copy connection conn
				 [in (bump-sequence-number in-state bytes-read)]))))
	(if (positive? remaining-to-read)
	    (let ((trailing-blocks (read-bytes/decrypt remaining-to-read in
						       (intra-packet-timeout) decryptor)))
	      (cond
	       ((false? trailing-blocks)
		(disconnect-with-error SSH_DISCONNECT_CONNECTION_LOST
				       "Timeout partway through reading a packet"))
	       ((eof-object? trailing-blocks)
		(if error-on-eof-or-timeout
		    (error 'read-packet "End-of-file during a packet")
		    (values trailing-blocks conn)))
	       (else
		(read-packet-trailer (bytes-append first-block trailing-blocks)))))
	    (read-packet-trailer first-block)))))))

;; 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)

;; DecodedPacket ConnectionState Optional<Boolean> -> ConnectionState
;; Encodes and writes a DecodedPacket to the ConnectionState.
(define (write-message! message conn [flush #f])
  (write `(sending ,message at out seq num ,(stream-state-sequence-number (connection-out conn))))
  (newline)
  (flush-output)
  (write-packet! (ssh-message-encode message) conn flush))

;; Bytes ConnectionState Boolean -> ConnectionState
;; Encrypts, MACs and writes a blob to the StreamState.
(define (write-packet! payload conn flush)
  (define out-state (connection-out conn))
  (define cipher (stream-state-cipher out-state))
  (define pad-block-size (supported-cipher-block-size (stream-state-cipher-description out-state)))
  (define out (stream-state-port out-state))
  (define encryptor (if cipher cipher values))
  ;; 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)
			      ((random-bytes padding-length) :: binary))))
  (define encrypted-packet (encryptor packet))
  (define computed-hmac-bytes (apply-hmac (stream-state-hmac out-state)
					  (stream-state-sequence-number out-state)
					  packet))
  (define mac-byte-count (bytes-length computed-hmac-bytes))
  (write-bytes encrypted-packet out)
  (when (positive? mac-byte-count)
    (write-bytes computed-hmac-bytes out))
  (when flush
    (flush-output out))
  (struct-copy connection conn
	       [out (bump-sequence-number out-state
					  (+ (bytes-length encrypted-packet) mac-byte-count))]))

;; 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))))

;; Bytes StreamState -> Natural
;; Reads and checks an HMAC for a received packet against its argument.
;; TODO:: Should the read HMAC bytes count against bytes-transferred?
(define (check-hmac! computed-hmac-bytes in)
  (define mac-byte-count (bytes-length computed-hmac-bytes))
  (when (positive? mac-byte-count)
    (let ((received-hmac-bytes (read-bytes/timeout mac-byte-count in
						   (intra-packet-timeout))))
      (cond
       ((false? received-hmac-bytes)
	(disconnect-with-error SSH_DISCONNECT_CONNECTION_LOST
			       "Timeout reading MAC"))
       ((eof-object? received-hmac-bytes)
	(disconnect-with-error SSH_DISCONNECT_CONNECTION_LOST
			       "EOF instead of MAC"))
       (else
	(when (not (equal? computed-hmac-bytes received-hmac-bytes))
	  (disconnect-with-error/local-info `((expected-hmac ,computed-hmac-bytes)
					      (actual-hmac ,received-hmac-bytes))
					    SSH_DISCONNECT_MAC_ERROR
					    "Corrupt MAC"))))))
  mac-byte-count)

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Packet dispatch and handling
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;; Bytes -> Byte
;; Retrieves the packet type byte from a packet.
(define (encoded-packet-msg-type encoded-packet)
  (bytes-ref encoded-packet 0))

;; PacketDispatcher [ Byte Maybe<PacketHandler> ]* -> PacketDispatcher
;; Adds or removes handlers to or from the given PacketDispatcher.
(define (extend-packet-dispatcher core-dispatcher . key-value-pairs)
  (let loop ((d core-dispatcher)
	     (key-value-pairs key-value-pairs))
    (cond
     ((null? key-value-pairs)
      d)
     ((null? (cdr key-value-pairs))
      (error 'extend-packet-dispatcher
	     "Must call extend-packet-dispatcher with matched key/value pairs"))
     (else
      (loop (let ((packet-type-number (car key-value-pairs))
		  (packet-handler-or-false (cadr key-value-pairs)))
	      (if packet-handler-or-false
		  (hash-set d packet-type-number packet-handler-or-false)
		  (hash-remove d packet-type-number)))
	    (cddr key-value-pairs))))))

;; ConnectionState [ Byte Maybe<PacketHandler> ]* -> ConnectionState
;; Installs (or removes) PacketHandlers in the given connection state;
;; see extend-packet-dispatcher.
(define (set-handlers conn . key-value-pairs)
  (struct-copy connection conn
	       [dispatch-table (apply extend-packet-dispatcher
				      (connection-dispatch-table conn)
				      key-value-pairs)]))

;; ConnectionState Byte PacketHandler -> ConnectionState
;; Installs a PacketHandler that removes the installed dispatch entry
;; and then delegates to its argument.
(define (oneshot-handler conn packet-type-number packet-handler)
  (set-handlers conn
		packet-type-number
		(lambda (packet message conn)
		  (packet-handler packet
				  message
				  (set-handlers conn packet-type-number #f)))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Handlers for core transport packet types
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;; PacketHandler for handling SSH_MSG_DISCONNECT.
(define (handle-msg-disconnect packet message conn)
  (raise (exn:fail:contract:protocol
	  (format "Received SSH_MSG_DISCONNECT with reason code ~a and message ~s"
		  (ssh-msg-disconnect-reason-code message)
		  (bytes->string/utf-8 (bit-string->bytes
					(ssh-msg-disconnect-description message))))
	  (current-continuation-marks)
	  (ssh-msg-disconnect-reason-code message)
	  '()
	  #t)))

;; PacketHandler for handling SSH_MSG_IGNORE.
(define (handle-msg-ignore packet message conn)
  ;; TODO: suppress debug printing.
  (write message)
  (newline)
  conn)

;; PacketHandler for handling SSH_MSG_UNIMPLEMENTED.
(define (handle-msg-unimplemented packet message conn)
  (disconnect-with-error/local-info
   `((offending-sequence-number ,(ssh-msg-unimplemented-sequence-number message)))
   SSH_DISCONNECT_PROTOCOL_ERROR
   "Disconnecting because of received SSH_MSG_UNIMPLEMENTED."))

;; PacketHandler for handling SSH_MSG_DEBUG.
(define (handle-msg-debug packet message conn)
  ;; TODO: use Racket log API.
  (write message)
  (newline)
  conn)

;; (SshMsgKexinit -> Symbol) SshMsgKexinit SshMsgKexinit -> Symbol
;; Computes the name of the "best" algorithm choice at the given
;; getter, using the rules from the RFC and the client and server
;; algorithm precedence lists.
(define (best-result getter client-algs server-algs)
  (define client-list0 (getter client-algs))
  (define server-list (getter server-algs))
  (let loop ((client-list client-list0))
    (cond
     ((null? client-list) (disconnect-with-error/local-info
			   `((client-list ,client-list0)
			     (server-list ,server-list))
			   SSH_DISCONNECT_KEY_EXCHANGE_FAILED
			   "Could not agree on a suitable algorithm for ~v"
			   getter))
     ((memq (car client-list) server-list) (car client-list))
     (else (loop (cdr client-list))))))

(define (check-host-key! host-key)
  ;; TODO: If we are *re*keying, worth checking here that the key hasn't *changed* either.
  (write `(TODO check-host-key! ,(hex (bit-string->bytes host-key)))) (newline) (flush-output)
  (void))

;; ExchangeHashInfo Bytes Natural Natural Natural -> Bytes
;; Computes the session ID as defined by SSH's DH key exchange method.
(define (dh-exchange-hash hash-info host-key e f k)
  (let ((block-to-hash
	 (bit-string->bytes
	  (bit-string ((string->bytes/utf-8 (exchange-hash-info-client-id hash-info)) :: (t:string))
		      ((string->bytes/utf-8 (exchange-hash-info-server-id hash-info)) :: (t:string))
		      ((exchange-hash-info-client-kexinit-bytes hash-info) :: (t:string))
		      ((exchange-hash-info-server-kexinit-bytes hash-info) :: (t:string))
		      (host-key :: (t:string))
		      (e :: (t:mpint))
		      (f :: (t:mpint))
		      (k :: (t:mpint))))))
    ;;(pretty-print `((block-to-hash ,(hex block-to-hash))))
    (sha1 block-to-hash)))

;; ExchangeHashInfo Symbol Symbol ConnectionState
;;   (Bytes Bytes Symbol ConnectionState -> ConnectionState)
;;   -> ConnectionState
;; Performs the client's half of the Diffie-Hellman key exchange protocol.
(define (perform-client-key-exchange hash-info kex-alg host-key-alg conn finish)
  (case kex-alg
    ((diffie-hellman-group14-sha1 diffie-hellman-group1-sha1)
     (define group (if (eq? kex-alg 'diffie-hellman-group14-sha1)
		       dh:oakley-group-14
		       dh:oakley-group-2)) ;; yes, SSH's group1 == Oakley/RFC2409 group 2
     (define-values (private-key public-key) (generate-key group))
     (oneshot-handler (write-message! (ssh-msg-kexdh-init (bit-string->integer public-key #t #f))
				      conn #t)
		      SSH_MSG_KEXDH_REPLY
		      (lambda (packet message conn)
			(define f (ssh-msg-kexdh-reply-f message))
			(define f-width (mpint-width f))
			(define f-as-bytes (integer->bit-string f (* 8 f-width) #t))
			(define shared-secret (compute-key private-key f-as-bytes))
			(define hash-alg sha1)
			(define exchange-hash
			  (dh-exchange-hash hash-info
					    (ssh-msg-kexdh-reply-host-key message)
					    (bit-string->integer public-key #t #f)
					    f
					    (bit-string->integer shared-secret #t #f)))
			;; (pretty-print `((public-key ,(hex public-key))
			;; 		(f-as-bytes ,(hex f-as-bytes))
			;; 		(shared-secret ,(hex shared-secret))
			;; 		(exchange-hash ,(hex exchange-hash))))
			(check-host-key! (ssh-msg-kexdh-reply-host-key message))
			(finish shared-secret exchange-hash hash-alg conn))))
    (else (disconnect-with-error SSH_DISCONNECT_KEY_EXCHANGE_FAILED
				 "Bad key-exchange algorithm ~v" kex-alg))))

;; StreamState Boolean Boolean (Bytes Maybe<Natural> -> Bytes)
;;   Symbol Symbol Symbol Symbol Symbol Symbol
;;   -> StreamState
;; Figures out which encryption, compression, and MAC option to use
;; for this stream, and initializes the relevant state vectors and
;; behaviours.
(define (apply-negotiated-options state is-server? is-outbound? derive-key
				  c2s-enc s2c-enc
				  c2s-mac s2c-mac
				  c2s-zip s2c-zip)
  ;; TODO: zip
  ;; TODO: make this less ugly. Compute all the keys, select just the ones we need afterward?
  (define c2s (if is-server? (not is-outbound?) is-outbound?)) ;; c2s true iff stream is serverward
  (define enc (if c2s c2s-enc s2c-enc))
  (define mac (if c2s c2s-mac s2c-mac))
  (define zip (if c2s c2s-zip s2c-zip))
  (struct-copy stream-state state
   [cipher (cond
	    ((assq enc supported-crypto-algorithms) =>
	     (lambda (entry)
	       (define c (cadr entry))
	       (define key (derive-key (if c2s #"C" #"D")
				       (supported-cipher-key-length c)))
	       (define iv (derive-key (if c2s #"A" #"B")
				      (supported-cipher-iv-length c)))
	       (define factory (supported-cipher-factory c))
	       ;; (pretty-print `(,is-server? ,(if c2s 'c2s 's2c) ,enc
	       ;; 				   (key ,(hex key)) (iv ,(hex iv))))
	       (factory is-outbound? key iv)))
	    (else (disconnect-with-error SSH_DISCONNECT_KEY_EXCHANGE_FAILED
					 "Could not find driver for encryption algorithm ~v"
					 enc)))]
   [hmac (cond
	  ((assq mac supported-hmac-algorithms) =>
	   (lambda (entry)
	     (define h (cadr entry))
	     (define factory (supported-hmac-factory h))
	     (define key (derive-key (if c2s #"E" #"F")
				     (supported-hmac-key-length h)))
	     ;; (pretty-print `(,is-server? ,(if c2s 'c2s 's2c) ,mac
	     ;; 				 (key ,(hex key))))
	     (factory key)))
	  (else (disconnect-with-error SSH_DISCONNECT_KEY_EXCHANGE_FAILED
				       "Could not find driver for HMAC algorithm ~v"
				       mac)))]))

(define (QQ conn)
  (write `(QQ ,(stream-state-sequence-number (connection-out conn)))) (newline)
  (write-message! (ssh-msg-debug #t #"Debug trace" #"") conn #t))

;; PacketHandler for handling SSH_MSG_KEXINIT.
(define (handle-msg-kexinit packet message conn)
  (define rekey (connection-rekey-state conn))
  (when (rekey-in-progress? rekey)
    (disconnect-with-error SSH_DISCONNECT_PROTOCOL_ERROR
			   "Received SSH_MSG_KEXINIT during ongoing key exchange"))
  (define local-algs (if (rekey-local? rekey)
			 (rekey-local-local-algorithms rekey)
			 ((local-algorithm-list))))
  (define encoded-local-algs (ssh-message-encode local-algs))
  (define remote-algs message)
  (define encoded-remote-algs packet)

  (when (rekey-wait? rekey)
    (set! conn (write-packet! encoded-local-algs conn #t)))

  (define is-server? (connection-is-server? conn))
  (define c (if is-server? remote-algs local-algs))
  (define s (if is-server? local-algs remote-algs))

  (define kex-alg (best-result ssh-msg-kexinit-kex_algorithms c s))
  (define host-key-alg (best-result ssh-msg-kexinit-server_host_key_algorithms c s))
  (define c2s-enc (best-result ssh-msg-kexinit-encryption_algorithms_client_to_server c s))
  (define s2c-enc (best-result ssh-msg-kexinit-encryption_algorithms_server_to_client c s))
  (define c2s-mac (best-result ssh-msg-kexinit-mac_algorithms_client_to_server c s))
  (define s2c-mac (best-result ssh-msg-kexinit-mac_algorithms_server_to_client c s))
  (define c2s-zip (best-result ssh-msg-kexinit-compression_algorithms_client_to_server c s))
  (define s2c-zip (best-result ssh-msg-kexinit-compression_algorithms_server_to_client c s))
  ;; Ignore languages.
  ;; Don't check the reserved field here, either. TODO: should we?

  (define (guess-matches? chosen-value getter)
    (let ((remote-choices (getter remote-algs)))
      (and (pair? remote-choices) ;; not strictly necessary because of
				  ;; the error behaviour of
				  ;; best-result.
	   (eq? (car remote-choices) ;; the remote peer's guess for this parameter
		chosen-value))))

  (define should-discard-first-kex-packet
    (and (ssh-msg-kexinit-first_kex_packet_follows remote-algs)
	 ;; They've already transmitted their guess. Does their guess match
	 ;; what we've actually selected?
	 (not (and
	       (guess-matches? kex-alg ssh-msg-kexinit-kex_algorithms)
	       (guess-matches? host-key-alg ssh-msg-kexinit-server_host_key_algorithms)
	       (guess-matches? c2s-enc ssh-msg-kexinit-encryption_algorithms_client_to_server)
	       (guess-matches? s2c-enc ssh-msg-kexinit-encryption_algorithms_server_to_client)
	       (guess-matches? c2s-mac ssh-msg-kexinit-mac_algorithms_client_to_server)
	       (guess-matches? s2c-mac ssh-msg-kexinit-mac_algorithms_server_to_client)
	       (guess-matches? c2s-zip ssh-msg-kexinit-compression_algorithms_client_to_server)
	       (guess-matches? s2c-zip ssh-msg-kexinit-compression_algorithms_server_to_client)))))

  (define (continue-after-discard conn)
    (if is-server?
	(error 'start-key-exchange "Server role unimplemented")
	(perform-client-key-exchange (exchange-hash-info (connection-local-id conn)
							 (connection-remote-id conn)
							 encoded-local-algs
							 encoded-remote-algs)
				     kex-alg
				     host-key-alg
				     conn
				     continue-after-key-exchange)))

  (define (continue-after-key-exchange shared-secret exchange-hash hash-alg conn)
    (define session-id (if (connection-session-id conn)
			   (connection-session-id conn) ;; don't overwrite existing ID
			   exchange-hash))
    (define k-h-prefix (bit-string ((bit-string->integer shared-secret #t #f) :: (t:mpint))
				   (exchange-hash :: binary)))
    (define (derive-key kind needed-bytes-or-false)
      (let extend ((key (hash-alg (bit-string->bytes
				   (bit-string (k-h-prefix :: binary)
					       (kind :: binary)
					       (session-id :: binary))))))
	(cond
	 ((false? needed-bytes-or-false)
	  key)
	 ((>= (bytes-length key) needed-bytes-or-false)
	  (subbytes key 0 needed-bytes-or-false))
	 (else
	  (extend (bytes-append key (hash-alg (bit-string->bytes
					       (bit-string (k-h-prefix :: binary)
							   (key :: binary))))))))))
    (oneshot-handler (struct-copy connection conn
				  [session-id session-id]) ;; just in case it changed
		     SSH_MSG_NEWKEYS
		     (lambda (newkeys-packet newkeys-message pre-newkeys-conn)
		       ;; First, send our SSH_MSG_NEWKEYS,
		       ;; incrementing the various counters, and then
		       ;; apply the new algorithms.
		       (define conn (write-message! (ssh-msg-newkeys) pre-newkeys-conn #t))
		       (struct-copy connection conn
				    [in 
				     (apply-negotiated-options (connection-in conn)
							       (connection-is-server? conn)
							       #f
							       derive-key
							       c2s-enc s2c-enc
							       c2s-mac s2c-mac
							       c2s-zip s2c-zip)]
				    [out
				     (apply-negotiated-options (connection-out conn)
							       (connection-is-server? conn)
							       #t
							       derive-key
							       c2s-enc s2c-enc
							       c2s-mac s2c-mac
							       c2s-zip s2c-zip)]
				    [rekey-state
				     (rekey-in-seconds-or-bytes (rekey-interval)
								(rekey-volume)
								(connection-in conn)
								(connection-out conn))]))))

  (if should-discard-first-kex-packet
      (let-values (((discarded-packet discarded-message conn) (read-message conn)))
	(continue-after-discard conn))
      (continue-after-discard conn)))

;; PacketDispatcher. Handles the core transport message types.
(define base-packet-dispatcher
  (hasheq SSH_MSG_DISCONNECT handle-msg-disconnect
	  SSH_MSG_IGNORE handle-msg-ignore
	  SSH_MSG_UNIMPLEMENTED handle-msg-unimplemented
	  SSH_MSG_DEBUG handle-msg-debug
	  SSH_MSG_KEXINIT handle-msg-kexinit))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Session choreography
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define (default-stream-state port)
  (stream-state port
		#f ;; cipher
		null-cipher-description
		0
		0
		null-hmac
		null-hmac-description
		(default-packet-limit)))

(define (send-preamble-and-identification! out)
  (let ((my-id (client-identification-string)))
    (for-each (lambda (line)
		(when (identification-line? line)
		  (error 'ssh-session
			 "Client preamble includes forbidden line ~v"
			 line))
		(display line out)
		(display "\r\n" out))
	      (client-preamble-lines))
    (display my-id out)
    (display "\r\n" out)
    (flush-output out)
    my-id))

;; Port -> String
(define (read-preamble-and-identification! in)
  (let ((line (read-line-limited in 253))) ;; 255 incl CRLF
    (when (eof-object? line)
      (error 'ssh-session "EOF while reading connection preamble"))
    (if (identification-line? line)
	line
	(read-preamble-and-identification! in))))

(define (ssh-session role in out)
  (define local-identification-string (send-preamble-and-identification! out))
  (with-handlers
      ((exn:fail? (lambda (e)
		    (close-input-port in)
		    (close-output-port out)
		    (raise e))))
    (let ((peer-identification-string (read-preamble-and-identification! in)))
      ;; Each identification string is both a cleartext indicator that
      ;; we've reached some notion of the right place and also input to
      ;; the hash function used during D-H key exchange.
      (when (not (regexp-match (required-peer-identification-regex)
			       peer-identification-string))
	(display "Invalid identification\r\n" out)
	(flush-output out)
	(error 'ssh-session
	       "Invalid peer identification string ~v"
	       peer-identification-string))

      (define result
	(let ((in-state (default-stream-state in))
	      (out-state (default-stream-state out)))
	  (run-ssh-session (connection in-state
				       out-state
				       base-packet-dispatcher
				       (hash) ;; TODO: make customizable
				       (lambda args
					 (error 'TODO-channel-open-handler))
				       (rekey-in-seconds-or-bytes -1 -1 in-state out-state)
				       (case role
					 ((client) #f)
					 ((server) #t))
				       local-identification-string
				       peer-identification-string
				       #f))))
      (close-input-port in)
      (close-output-port out)
      result)))

(define (rekey-in-seconds-or-bytes delta-seconds delta-bytes in-state out-state)
  (rekey-wait (+ (current-seconds) delta-seconds)
	      (+ (stream-state-bytes-transferred in-state)
		 (stream-state-bytes-transferred out-state)
		 delta-bytes)))

(define (time-to-rekey? rekey conn)
  (and (rekey-wait? rekey)
       (or (>= (current-seconds) (rekey-wait-deadline rekey))
	   (>= (+ (stream-state-bytes-transferred (connection-in conn))
		  (stream-state-bytes-transferred (connection-out conn)))
	       (rekey-wait-threshold-bytes rekey)))))

(define (maybe-send-disconnect-message! e conn)
  (if (exn:fail:contract:protocol-originated-at-peer? e)
      conn
      (write-message! (ssh-msg-disconnect (exn:fail:contract:protocol-reason-code e)
					  (string->bytes/utf-8 (exn-message e))
					  #"")
		      conn
		      #t)))

(define (write-messages! outbound-messages conn)
  (let ((final-state (foldl write-message! conn outbound-messages)))
    (flush-output (stream-state-port (connection-out final-state)))
    final-state))

;; ConnectionState -> TODO:?
(define (run-ssh-session conn)
  (with-handlers
      ((exn:fail:contract:protocol? (lambda (e)
				      (maybe-send-disconnect-message! e conn)
				      (raise e))))
    (let loop ((new-connection-state conn))
      ;; YUCK: in order to be able to send our disconnect messages in
      ;; the with-handlers above, we need to know the most up-to-date
      ;; connection state. This is a thorny, ugly problem.
      (set! conn new-connection-state)
      (if (time-to-rekey? (connection-rekey-state conn) conn)
	  (let ((algs ((local-algorithm-list))))
	    (loop (struct-copy connection (write-message! algs conn #t)
			       [rekey-state (rekey-local algs)])))
	  (let-values (((packet message conn) (read-message conn #f)))
	    (cond
	     ((eof-object? packet)
	      (error 'TODO-disconnected-without-shutdown))
	     ((false? packet)
	      ;; Timeout waiting for a message.
	      (loop conn))
	     (else
	      (let* ((packet-type-number (encoded-packet-msg-type packet))
		     (packet-handler (hash-ref (connection-dispatch-table conn)
					       packet-type-number
					       #f)))
		(if packet-handler
		    (loop (packet-handler packet message conn))
		    (loop (ssh-msg-unimplemented
				       (most-recent-received-sequence-number conn))))))))))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Test driver code
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(require racket/tcp)
(require racket/pretty)
(define (t)
  (let-values (((i o) (tcp-connect "localhost"
				   2323
				   ;;22
				   )))
    (ssh-session 'client i o)))
(t)