#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")

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

(struct stream-state (port
		      cipher
		      cipher-description
		      sequence-number
		      bytes-transferred
		      hmac
		      hmac-description
		      packet-size-limit)
	#: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)

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

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

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

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

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

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

(define (bump-sequence-number! state byte-count)
  (struct-copy stream-state state
	       [sequence-number
		(+ 1 (stream-state-sequence-number state))]
	       [bytes-transferred
		(+ byte-count (stream-state-bytes-transferred state))]))

;; 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 for end-of-file or
;; timeout, respectively, depending on error-on-eof-or-timeout.
(define (read-message in-state out-state [error-on-eof-or-timeout #t])
  (let-values (((packet in-state) (read-packet in-state error-on-eof-or-timeout)))
    (if (not (bytes? packet))
	(values packet packet in-state out-state)
	(let ((message (ssh-message-decode packet)))
	  (write `(received ,message)) (newline) (flush-output)
	  (if message
	      (values packet message in-state out-state)
	      (let ((bad-seq-num (- (stream-state-sequence-number in-state) 1)))
		;; TODO: remove this debug output
		(display "BAD PACKET ")
		(display (hex packet))
		(newline)
		(flush-output)
		(read-message in-state
			      (write-message! (ssh-msg-unimplemented bad-seq-num)
					      out-state)
			      error-on-eof-or-timeout)))))))

;; 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
(define (read-packet in-state error-on-eof-or-timeout)
  (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 in-state)))
   ((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 in-state)))
   (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))
		    (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 in-state)))
	       (else
		(read-packet-trailer (bytes-append first-block trailing-blocks)))))
	    (read-packet-trailer first-block)))))))

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

(define (write-message! message out-state [flush #f])
  (write `(sending ,message)) (newline) (flush-output)
  (write-packet! (ssh-message-encode message) out-state flush))

(define (write-packet! payload out-state flush)
  (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 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))
  (bump-sequence-number! out-state (+ (bytes-length encrypted-packet) mac-byte-count)))

(define (apply-hmac mac sequence-number packet)
  (mac (bit-string->bytes (bit-string-append (integer->bit-string sequence-number 32 #t) packet))))

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

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

(define (null-hmac blob)
  #"")

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

(define (encoded-packet-msg-type encoded-packet)
  (bytes-ref encoded-packet 0))

(define (key-exchange-init? encoded-packet)
  (= (encoded-packet-msg-type encoded-packet) SSH_MSG_KEXINIT))

(define (acceptable-during-key-exchange? encoded-packet)
  ;; See end of RFC 4253 section 7.1.
  (let ((msg-type (encoded-packet-msg-type encoded-packet)))
    (and (ssh-msg-type-transport-layer? msg-type)
	 (not (memv msg-type (list SSH_MSG_SERVICE_REQUEST
				   SSH_MSG_SERVICE_ACCEPT
				   SSH_MSG_KEXINIT))))))

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

(define (ssh-session role in out seed message-handler)
  (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 ((orig-in-state (default-stream-state in))
	      (orig-out-state (default-stream-state out)))
	  (ssh-session-loop role
			    local-identification-string
			    peer-identification-string
			    message-handler
			    #f
			    seed
			    (rekey-in-seconds-or-bytes -1 -1 orig-in-state orig-out-state)
			    orig-in-state
			    orig-out-state)))

      (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 in-state out-state)
  (and (rekey-wait? rekey)
       (or (>= (current-seconds) (rekey-wait-deadline rekey))
	   (>= (+ (stream-state-bytes-transferred in-state)
		  (stream-state-bytes-transferred out-state))
	       (rekey-wait-threshold-bytes rekey)))))

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

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

(define (ssh-session-loop role local-id remote-id message-handler
			  session-id
			  seed rekey in-state out-state)
  (let loop ((seed seed)
	     (rekey rekey)
	     (in-state in-state)
	     (out-state out-state))
    (with-handlers
	((exn:fail:contract:protocol? (lambda (e)
					(maybe-send-disconnect-message! e out-state)
					(raise e))))
      (if (time-to-rekey? rekey in-state out-state)
	  (let ((algs ((local-algorithm-list))))
	    (loop seed
		  (rekey-local algs)
		  in-state
		  (write-message! algs out-state #t)))
	  (let-values (((packet message in-state out-state)
			(read-message in-state out-state #f)))
	    (cond
	     ((eof-object? packet) seed)
	     ((false? packet)
	      ;; Timeout waiting for a message.
	      (loop seed rekey in-state out-state))
	     ((key-exchange-init? packet)
	      (if (rekey-in-progress? rekey)
		  (disconnect-with-error SSH_DISCONNECT_PROTOCOL_ERROR
					 "Received SSH_MSG_KEXINIT during ongoing key exchange")
		  (let* ((algs (if (rekey-local? rekey)
				   (rekey-local-local-algorithms rekey)
				   ((local-algorithm-list))))
			 (encoded-algs (ssh-message-encode algs))
			 (out-state (if (rekey-wait? rekey)
					(write-packet! encoded-algs out-state #t)
					out-state)))
		    (start-key-exchange session-id ;; may be #f, in which case will change below
					role
					local-id
					encoded-algs
					algs
					remote-id
					packet
					message
					in-state
					out-state
					(lambda (session-id in-state out-state)
					  (ssh-session-loop
					   role
					   local-id
					   remote-id
					   message-handler
					   session-id ;; just in case it changed
					   seed
					   (rekey-in-seconds-or-bytes (rekey-interval)
								      (rekey-volume)
								      in-state
								      out-state)
					   in-state
					   (write-message! (ssh-msg-ignore #"hello world")
							   out-state #t)))))))
	     ((and (rekey-in-progress? rekey)
		   (not (acceptable-during-key-exchange? packet)))
	      (disconnect-with-error SSH_DISCONNECT_PROTOCOL_ERROR
				     "Unacceptable packet type ~v"
				     (encoded-packet-msg-type packet)))
	     (else
	      (message-handler seed
			       message
			       (lambda (outbound-messages new-seed)
				 (loop new-seed
				       rekey
				       in-state
				       (write-messages! outbound-messages out-state)))))))))))

(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 (apply-negotiated-options state role 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 (case role
		((client) outbound)
		((server) (not 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))
  (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 `(,role ,(if c2s 'c2s 's2c) (key ,(hex key)) (iv ,(hex iv))))
	       (factory 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 `(,role ,(if c2s 'c2s 's2c) (key ,(hex key))))
	     (factory key)))
	  (else (disconnect-with-error SSH_DISCONNECT_KEY_EXCHANGE_FAILED
				       "Could not find driver for HMAC algorithm ~v"
				       mac)))]))

(define (start-key-exchange old-session-id ;; bytes or #f
			    role ;; 'client or 'server
			    local-id ;; string
			    encoded-local-algs ;; bytes, an encoded ssh-msg-kexinit
			    local-algs ;; ssh-msg-kexinit
			    remote-id ;; string
			    encoded-remote-algs ;; bytes, an encoded ssh-msg-kexinit
			    remote-algs ;; ssh-msg-kexinit
			    in-state
			    out-state
			    finish-key-exchange)
  (when (not (memq role '(client server)))
    (error 'start-key-exchange "Illegal role ~v, must be either 'client or 'server" role))

  (define c (case role ((client) local-algs) ((server) remote-algs)))
  (define s (case role ((client) remote-algs) ((server) local-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 in-state out-state)
    (case role
      ((client) (perform-client-key-exchange (exchange-hash-info local-id
								 remote-id
								 encoded-local-algs
								 encoded-remote-algs)
					     kex-alg host-key-alg in-state out-state
					     continue-after-key-exchange))
      ((server) (error 'start-key-exchange "Server role unimplemented"))))

  (define (continue-after-key-exchange shared-secret exchange-hash hash-alg in-state out-state)
    (let ((session-id (if old-session-id
			  old-session-id ;; don't overwrite existing ID
			  exchange-hash))
	  (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))))))))))
      (let-values (((newkeys-packet newkeys-message in-state out-state)
		    (read-message in-state out-state)))
	(when (not (ssh-msg-newkeys? newkeys-message))
	  (disconnect-with-error/local-info `((message ,newkeys-message))
					    SSH_DISCONNECT_PROTOCOL_ERROR
					    "Expected SSH_MSG_NEWKEYS"))
	(let ((out-state (write-message! (ssh-msg-newkeys) out-state)))
	  (finish-key-exchange session-id
			       (apply-negotiated-options in-state role #f derive-key
							 c2s-enc s2c-enc
							 c2s-mac s2c-mac
							 c2s-zip s2c-zip)
			       (apply-negotiated-options out-state role #t derive-key
							 c2s-enc s2c-enc
							 c2s-mac s2c-mac
							 c2s-zip s2c-zip))))))

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

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

(define (perform-client-key-exchange hash-info kex-alg host-key-alg in-state out-state finish)
  (case kex-alg
    ((diffie-hellman-group14-sha1 diffie-hellman-group1-sha1)
     (let ((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
       (let*-values (((private-key public-key) (generate-key group))
		     ((out-state)
		      (write-message! (ssh-msg-kexdh-init (bit-string->integer public-key #t #f))
				      out-state
				      #t))
		     ((packet message in-state out-state)
		      (read-message in-state out-state)))
	 (if (not (ssh-msg-kexdh-reply? message))
	     (disconnect-with-error/local-info `((message ,message))
					       SSH_DISCONNECT_KEY_EXCHANGE_FAILED
					       "Unexpected packet type")
	     (let* ((f (ssh-msg-kexdh-reply-f message))
		    (f-width (mpint-width f))
		    (f-as-bytes (integer->bit-string f (* 8 f-width) #t))
		    (shared-secret (compute-key private-key f-as-bytes))
		    (hash-alg sha1)
		    (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 in-state out-state))))))
    (else (disconnect-with-error SSH_DISCONNECT_KEY_EXCHANGE_FAILED
				 "Bad key-exchange algorithm ~v" kex-alg))))

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

(require racket/tcp)
(require racket/pretty)
(define (t)
  (let-values (((i o) (tcp-connect "localhost"
				   2323
				   ;;22
				   )))
    (ssh-session 'client
		 i o #f
		 (lambda (seed message continue-reading)
		   (pretty-print message)
		   (continue-reading (list) seed)))))
(t)