#lang racket/base
;;
;;; Copyright 2010, 2011, 2012, 2013 Tony Garnock-Jones <tonyg@ccs.neu.edu>
;;;
;;; This file is part of marketplace-ssh.
;;;
;;; marketplace-ssh is free software: you can redistribute it and/or
;;; modify it under the terms of the GNU General Public License as
;;; published by the Free Software Foundation, either version 3 of the
;;; License, or (at your option) any later version.
;;;
;;; marketplace-ssh is distributed in the hope that it will be useful,
;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
;;; General Public License for more details.
;;;
;;; You should have received a copy of the GNU General Public License
;;; along with marketplace-ssh. If not, see
;;; <http://www.gnu.org/licenses/>.

(require "ssh-numbers.rkt")

(require (for-syntax racket/base))
(require (for-syntax (only-in racket/list append*)))
(require (for-syntax (only-in srfi/1 iota)))

(require bitsyntax)
(require racket/bytes)

(require rackunit)

(provide ssh-message-decode
	 ssh-message-encode)

(provide t:boolean
	 t:string
	 t:mpint
	 mpint-width
	 t:name-list)

(provide (struct-out ssh-msg-kexinit)
	 (struct-out ssh-msg-kexdh-init)
	 (struct-out ssh-msg-kexdh-reply)
	 (struct-out ssh-msg-disconnect)
	 (struct-out ssh-msg-unimplemented)
	 (struct-out ssh-msg-newkeys)
	 (struct-out ssh-msg-debug)
	 (struct-out ssh-msg-ignore)
	 (struct-out ssh-msg-service-request)
	 (struct-out ssh-msg-service-accept)
	 (struct-out ssh-msg-userauth-request)
	 (struct-out ssh-msg-userauth-failure)
	 (struct-out ssh-msg-userauth-success)
	 (struct-out ssh-msg-global-request)
	 (struct-out ssh-msg-request-success)
	 (struct-out ssh-msg-request-failure)
	 (struct-out ssh-msg-channel-open)
	 (struct-out ssh-msg-channel-open-confirmation)
	 (struct-out ssh-msg-channel-open-failure)
	 (struct-out ssh-msg-channel-window-adjust)
	 (struct-out ssh-msg-channel-data)
	 (struct-out ssh-msg-channel-extended-data)
	 (struct-out ssh-msg-channel-eof)
	 (struct-out ssh-msg-channel-close)
	 (struct-out ssh-msg-channel-request)
	 (struct-out ssh-msg-channel-success)
	 (struct-out ssh-msg-channel-failure)
	 )

(define encoder-map (make-hasheqv))
(define decoder-map (make-hasheqv))

(define (ssh-message-decode packet)
  (define type-code (bytes-ref packet 0))
  (define decoder (hash-ref decoder-map type-code #f))
  (if decoder
      (decoder packet)
      #f))

(define (ssh-message-encode m)
  (bit-string->bytes ((hash-ref encoder-map (prefab-struct-key m)) m)))

(define-syntax define-ssh-message-type
  (syntax-rules ()
    ((_ name type-byte-value (field-type field-name) ...)
     (begin
       (struct name (field-name ...) #:prefab)
       (hash-set! encoder-map 'name
		  (compute-ssh-message-encoder type-byte-value field-type ...))
       (hash-set! decoder-map type-byte-value
		  (compute-ssh-message-decoder name type-byte-value field-type ...))))))

(define-syntax t:boolean
  (syntax-rules ()
    ((_ #t input ks kf)
     (bit-string-case input
       ([ v (rest :: binary) ]
	(ks (not (zero? v)) rest))
       (else (kf))))
    ((_ #f v) (bit-string (if v 1 0)))))

(define-syntax t:packed-bytes
  (syntax-rules ()
    ((_ #t input ks kf n)
     (bit-string-case input
       ([ (bs :: binary bytes n) (rest :: binary) ]
	(ks (bit-string->bytes bs) rest))
       (else (kf))))
    ((_ #t input ks kf)
     (bit-string-case input
       ([ (rest :: binary) ]
	(ks (bit-string->bytes rest) #""))
       (else (kf))))
    ((_ #f bs n) (bit-string (bs :: binary)))
    ((_ #f bs) (bit-string (bs :: binary)))))

(define-syntax t:string
  (syntax-rules ()
    ((_ #t input ks kf #:pack)
     (t:string #t input (lambda (v rest) (ks (bit-string->bytes v) rest)) kf))
    ((_ #t input ks kf)
     (bit-string-case input
       ([ (length :: integer bits 32) (body :: binary bytes length) (rest :: binary) ]
	(ks body rest))
       (else (kf))))
    ((_ #f bs #:pack) (t:string #f bs)) ;; #:pack ignored on encoding
    ((_ #f bs)
     (bit-string ((bytes-length (bit-string->bytes bs)) :: integer bits 32)
		 (bs :: binary)))))

(define-syntax t:mpint
  (syntax-rules ()
    ((_ #t input ks kf)
     (bit-string-case input
       ([ (length :: integer bits 32) (body :: binary bytes length) (rest :: binary) ]
	(ks (if (zero? (bit-string-length body)) 0 (bit-string->integer body #t #t))
	    rest))
       (else (kf))))
    ((_ #f n)
     (let* ((width (mpint-width n))
	    (buf (integer->bit-string n (* 8 width) #t)))
       (bit-string (width :: integer bits 32) (buf :: binary))))))

(define-syntax t:name-list
  (syntax-rules ()
    ((_ #t input ks kf)
     (t:string #t
	       input
	       (lambda (body rest) (ks (name-list->symbols body) rest))
	       kf))
    ((_ #f ns)
     (t:string #f (symbols->name-list ns)))))

(define-for-syntax (codec-options field-type)
  (syntax-case field-type (byte boolean uint32 uint64 string mpint name-list)
    (byte #'(integer bits 8))
    ((byte n) #'((t:packed-bytes n)))
    (boolean #'((t:boolean)))
    (uint32 #'(integer bits 32))
    (uint64 #'(integer bits 64))
    (string #'((t:string #:pack)))
    (mpint #'((t:mpint)))
    (name-list #'((t:name-list)))
    (extension #'((t:packed-bytes)))))

(define-syntax compute-ssh-message-encoder
  (lambda (stx)
    (syntax-case stx ()
      ((_ type-byte-value field-type ...)
       #`(lambda (message)
	   (let ((vec (struct->vector message)))
	     #,(with-syntax (((field-spec ...)
			      (let ((type-list (syntax->list #'(field-type ...))))
				(map (lambda (index type)
				       #`((vector-ref vec #,index) :: #,@(codec-options type)))
				     (iota (length type-list) 1)
				     type-list))))
		 #'(bit-string (type-byte-value :: integer bytes 1)
			       field-spec ...))))))))

(define-syntax compute-ssh-message-decoder
  (lambda (stx)
    (syntax-case stx ()
      ((_ struct-name type-byte-value field-type ...)
       (with-syntax (((temp-name ...) (generate-temporaries #'(field-type ...)))
		     (((codec-option ...) ...)
		      (map codec-options (syntax->list #'(field-type ...)))))
	 #`(lambda (packet)
	     (bit-string-case packet
	       ([ (= type-byte-value) (temp-name :: codec-option ...) ... ]
		(struct-name temp-name ...)))))))))

(define (mpint-width n)
  (if (zero? n)
      0
      (+ 1 (quotient (integer-length n) 8))))

(check-eqv? (mpint-width 0) 0)
(check-eqv? (mpint-width #x9a378f9b2e332a7) 8)
(check-eqv? (mpint-width #x7f) 1)
(check-eqv? (mpint-width #x80) 2)
(check-eqv? (mpint-width #x81) 2)
(check-eqv? (mpint-width #xff) 2)
(check-eqv? (mpint-width #x100) 2)
(check-eqv? (mpint-width #x101) 2)
(check-eqv? (mpint-width #x-1234) 2)
(check-eqv? (mpint-width #x-deadbeef) 5)

(define (symbols->name-list syms)
  (bytes-join (map (lambda (s) (string->bytes/utf-8 (symbol->string s))) syms) #","))

(define (name-list->symbols bs)
  (if (zero? (bit-string-length bs))
      '()
      (map string->symbol (regexp-split #rx"," (bytes->string/utf-8 (bit-string->bytes bs))))))

(struct test-message (value) #:prefab)
(let ((test-decode (compute-ssh-message-decoder test-message 123 mpint))
      (test-encode (compute-ssh-message-encoder 123 mpint)))
  (define (bidi-check msg enc-without-type-tag)
    (let ((enc (bytes-append (bytes 123) enc-without-type-tag)))
      (let ((msg-enc (bit-string->bytes (test-encode msg)))
	    (enc-msg (test-decode enc)))
	(if (and (equal? msg-enc enc)
		 (equal? enc-msg msg))
	    'ok
	    `(fail ,msg-enc ,enc-msg)))))
  (check-eqv? (bidi-check (test-message 0) (bytes 0 0 0 0)) 'ok)
  (check-eqv? (bidi-check (test-message #x9a378f9b2e332a7)
			  (bytes #x00 #x00 #x00 #x08
				 #x09 #xa3 #x78 #xf9
				 #xb2 #xe3 #x32 #xa7)) 'ok)
  (check-eqv? (bidi-check (test-message #x80)
			  (bytes #x00 #x00 #x00 #x02 #x00 #x80)) 'ok)
  (check-eqv? (bidi-check (test-message #x-1234)
			  (bytes #x00 #x00 #x00 #x02 #xed #xcc)) 'ok)
  (check-eqv? (bidi-check (test-message #x-deadbeef)
			  (bytes #x00 #x00 #x00 #x05
				 #xff #x21 #x52 #x41 #x11)) 'ok))

(define-ssh-message-type ssh-msg-kexinit SSH_MSG_KEXINIT
     ((byte 16)    cookie)
     (name-list    kex_algorithms)
     (name-list    server_host_key_algorithms)
     (name-list    encryption_algorithms_client_to_server)
     (name-list    encryption_algorithms_server_to_client)
     (name-list    mac_algorithms_client_to_server)
     (name-list    mac_algorithms_server_to_client)
     (name-list    compression_algorithms_client_to_server)
     (name-list    compression_algorithms_server_to_client)
     (name-list    languages_client_to_server)
     (name-list    languages_server_to_client)
     (boolean      first_kex_packet_follows)
     (uint32       reserved))

(define-ssh-message-type ssh-msg-kexdh-init SSH_MSG_KEXDH_INIT
  (mpint e))

(define-ssh-message-type ssh-msg-kexdh-reply SSH_MSG_KEXDH_REPLY
  (string host-key)
  (mpint f)
  (string h-signature))

(define-ssh-message-type ssh-msg-disconnect SSH_MSG_DISCONNECT
  (uint32 reason-code)
  (string description)
  ;; TODO: OpenSSH 5.3p1 Debian-3ubuntu7 25 Mar 2009 (from lucid)
  ;; sends SSH_MSG_DISCONNECT without the language-tag field! In
  ;; particular, when I press ^D to terminate my session, I get
  ;; #"\1\0\0\0\v\0\0\0\24disconnected by user".
  (string language-tag))

(define-ssh-message-type ssh-msg-unimplemented SSH_MSG_UNIMPLEMENTED
  (uint32 sequence-number))

(define-ssh-message-type ssh-msg-newkeys SSH_MSG_NEWKEYS)

(define-ssh-message-type ssh-msg-debug SSH_MSG_DEBUG
  (boolean always-display?)
  (string message)
  (string language-tag))

(define-ssh-message-type ssh-msg-ignore SSH_MSG_IGNORE
  (string data))

(define-ssh-message-type ssh-msg-service-request SSH_MSG_SERVICE_REQUEST
  (string service-name))

(define-ssh-message-type ssh-msg-service-accept SSH_MSG_SERVICE_ACCEPT
  (string service-name))

(define-ssh-message-type ssh-msg-userauth-request SSH_MSG_USERAUTH_REQUEST
  (string    user-name)
  (string    service-name)
  (string    method-name)
  (extension method-specific-fields))

(define-ssh-message-type ssh-msg-userauth-failure SSH_MSG_USERAUTH_FAILURE
  (name-list    continuable-authentications)
  (boolean      partial-success?))

(define-ssh-message-type ssh-msg-userauth-success SSH_MSG_USERAUTH_SUCCESS)

(define-ssh-message-type ssh-msg-global-request SSH_MSG_GLOBAL_REQUEST
  (string    request-name)
  (boolean   want-reply?)
  (extension data))

(define-ssh-message-type ssh-msg-request-success SSH_MSG_REQUEST_SUCCESS
  (extension data))

(define-ssh-message-type ssh-msg-request-failure SSH_MSG_REQUEST_FAILURE)

(define-ssh-message-type ssh-msg-channel-open SSH_MSG_CHANNEL_OPEN
  (string    channel-type)
  (uint32    sender-channel)
  (uint32    initial-window-size)
  (uint32    maximum-packet-size)
  (extension data))

(define-ssh-message-type ssh-msg-channel-open-confirmation SSH_MSG_CHANNEL_OPEN_CONFIRMATION
  (uint32    recipient-channel)
  (uint32    sender-channel)
  (uint32    initial-window-size)
  (uint32    maximum-packet-size)
  (extension data))

(define-ssh-message-type ssh-msg-channel-open-failure SSH_MSG_CHANNEL_OPEN_FAILURE
  (uint32    recipient-channel)
  (uint32    reason)
  (string    description)
  (string    language))

(define-ssh-message-type ssh-msg-channel-window-adjust SSH_MSG_CHANNEL_WINDOW_ADJUST
  (uint32    recipient-channel)
  (uint32    bytes))

(define-ssh-message-type ssh-msg-channel-data SSH_MSG_CHANNEL_DATA
  (uint32    recipient-channel)
  (string    data))

(define-ssh-message-type ssh-msg-channel-extended-data SSH_MSG_CHANNEL_EXTENDED_DATA
  (uint32    recipient-channel)
  (uint32    type-code)
  (string    data))

(define-ssh-message-type ssh-msg-channel-eof SSH_MSG_CHANNEL_EOF
  (uint32    recipient-channel))

(define-ssh-message-type ssh-msg-channel-close SSH_MSG_CHANNEL_CLOSE
  (uint32    recipient-channel))

(define-ssh-message-type ssh-msg-channel-request SSH_MSG_CHANNEL_REQUEST
  (uint32    recipient-channel)
  (string    type)
  (boolean   want-reply?)
  (extension data))

(define-ssh-message-type ssh-msg-channel-success SSH_MSG_CHANNEL_SUCCESS
  (uint32    recipient-channel))

(define-ssh-message-type ssh-msg-channel-failure SSH_MSG_CHANNEL_FAILURE
  (uint32    recipient-channel))