#lang racket/base

(require (planet tonyg/bitsyntax))
(require (planet vyzo/crypto:2:3))

(require rackunit)
(require "safe-io.rkt")
(require "ssh-numbers.rkt")
(require "ssh-message-types.rkt")

(define identification-recogniser #rx"^SSH-")
(define (identification-line? str)
  (regexp-match identification-recogniser str))

(define required-server-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)
  (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 (client-identification-string) out)
  (display "\r\n" out)
  (flush-output out))

(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 (check-packet-length! actual-length limit block-size)
  (when (> actual-length limit)
    (error 'check-packet-length!
	   "Packet of length ~v is longer than packet limit ~v"
	   actual-length
	   limit))
  (when (not (zero? (modulo actual-length block-size)))
    (error 'check-packet-length!
	   "Packet of length ~v is not a multiple of block size ~v"
	   actual-length
	   block-size)))

(define (read-bytes/decrypt count in decryptor)
  (let ((encrypted (read-bytes count in)))
    (if (eof-object? encrypted)
	encrypted
	(decryptor encrypted))))

;; 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
		     cipher
		     sequence-number
		     check-mac!
		     packet-size-limit)
  (define first-block-size (if cipher (cipher-block-size cipher) 8))
  (define subsequent-block-size (if cipher first-block-size 1))
  (define decryptor (if cipher (lambda (blocks) (cipher-update! cipher blocks)) values))
  (define first-block (read-bytes/decrypt first-block-size in decryptor))
  (if (eof-object? first-block)
      first-block
      (let* ((packet-length (integer-bytes->integer first-block #f #t 0 4)))
	(check-packet-length! packet-length packet-size-limit subsequent-block-size)
	(let* ((padding-length (bytes-ref first-block 4))
	       (payload-length (- packet-length padding-length 1))
	       (remaining-to-read (- packet-length (bytes-length first-block))))
	  (define (read-packet-trailer packet)
	    (check-mac! sequence-number packet in)
	    (subbytes packet 5 (+ 5 payload-length)))
	  (if (positive? remaining-to-read)
	      (let ((trailing-blocks (read-bytes/decrypt remaining-to-read in decryptor)))
		(if (eof-object? trailing-blocks)
		    trailing-blocks
		    (read-packet-trailer (bytes-append first-block trailing-blocks))))
	      (read-packet-trailer first-block))))))

(define (null-mac-checker sequence-number packet in)
  'ok)

(define (ssh-session in out seed packet-handler)
  (send-preamble-and-identification! out)
  (let ((server-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-server-identification-regex)
			     server-identification-string))
      (error 'ssh-session "Received invalid identification ~v from peer"
	     server-identification-string))

    (define (write-packet! . args)
      (error 'write-packet "Unimplemented"))

    (define result
      (let loop ((seed seed)
		 (sequence-number 0)
		 (cipher #f)
		 (check-mac! null-mac-checker))
	(let ((packet (read-packet in cipher sequence-number check-mac! 65536)))
	  (if (eof-object? packet)
	      seed
	      (packet-handler seed
			      packet
			      write-packet!
			      (lambda (new-seed) (loop new-seed
						       (bitwise-and (+ sequence-number 1)
								    #xffffffff)
						       cipher
						       check-mac!))
			      (lambda args
				(error 'rekey "Unimplemented")))))))

    (close-input-port in)
    (close-output-port out)
    result))

(require racket/tcp)
(require racket/pretty)
(define (t)
  (let-values (((i o) (tcp-connect "localhost" 22)))
    (ssh-session i o #f
		 (lambda (seed packet write-packet! continue-reading rekey)
		   (pretty-print packet)
		   (pretty-print (ssh-message-decode packet))
		   (continue-reading seed)))))
(t)