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

(require racket/match)
(require racket/port)
(require net/base64)

(require bitsyntax)
(require "crypto.rkt")
(require "asn1-ber.rkt")
(require "ssh-message-types.rkt")

(require rackunit)

(provide (struct-out rsa-private-key)
	 (struct-out dsa-private-key)
	 (struct-out rsa-public-key)
	 (struct-out dsa-public-key)

	 public-key->pieces
	 pieces->public-key

	 host-key-algorithm->keys
	 host-key-signature
	 verify-host-key-signature!

	 pieces->ssh-host-key
	 ssh-host-key->pieces)

(struct rsa-private-key (version n e d p q dmp1 dmq1 iqmp) #:transparent)
(struct dsa-private-key (version p q g y x) #:transparent)
(struct rsa-public-key (n e) #:transparent)
(struct dsa-public-key (y p q g) #:transparent)

;; ASN.1 BER integers are signed.
(define (bs->n bs) (bit-string->integer bs #t #t))
(define (n->bs n) (integer->bit-string n (* 8 (mpint-width n)) #t))

;; (define (private-key->pieces key)
;;   (bytes->private-key-pieces (private-key->bytes key)))

(define (bytes->private-key-pieces bs)
  (match (asn1-ber-decode-all bs)
    (`(0 16 ((0 2 ,version-bytes)
	     (0 2 ,n-bytes)
	     (0 2 ,e-bytes)
	     (0 2 ,d-bytes)
	     (0 2 ,p-bytes)
	     (0 2 ,q-bytes)
	     (0 2 ,dmp1-bytes)
	     (0 2 ,dmq1-bytes)
	     (0 2 ,iqmp-bytes)))
     (rsa-private-key (bs->n version-bytes)
		      (bs->n n-bytes)
		      (bs->n e-bytes)
		      (bs->n d-bytes)
		      (bs->n p-bytes)
		      (bs->n q-bytes)
		      (bs->n dmp1-bytes)
		      (bs->n dmq1-bytes)
		      (bs->n iqmp-bytes)))
    (`(0 16 ((0 2 ,version-bytes)
	     (0 2 ,p-bytes)
	     (0 2 ,q-bytes)
	     (0 2 ,g-bytes)
	     (0 2 ,public-key-bytes) ;; y
	     (0 2 ,private-key-bytes))) ;; x
     (dsa-private-key (bs->n version-bytes)
		      (bs->n p-bytes)
		      (bs->n q-bytes)
		      (bs->n g-bytes)
		      (bs->n public-key-bytes)
		      (bs->n private-key-bytes)))))

(define (pieces->private-key p)
  (match p
    [(struct rsa-private-key (version n e d _p _q _dmp1 _dmq1 _iqmp))
     (datum->pk-key (list 'rsa 'private n e d)
                    'rkt-private)]
    [(struct dsa-private-key (version p q g y x))
     (datum->pk-key (list 'dsa 'private p q g y x)
                    'rkt-private)]))

(define (public-key->pieces key)
  (match (pk-key->datum key 'rkt-public)
    [(list 'rsa 'public n e)
     (rsa-public-key n e)]
    [(list 'dsa 'public p q g public-key-bytes)
     (dsa-public-key public-key-bytes p q g)]))

(define (pieces->public-key p)
  (match p
    ((struct rsa-public-key (n e))
     (datum->pk-key (list 'rsa 'public n e) 'rkt-public))
    ((struct dsa-public-key (y p q g))
     (datum->pk-key (list 'dsa 'public p q g y) 'rkt-public))))

(define (host-key-algorithm->keys host-key-alg)
  (case host-key-alg
    ((ssh-dss) (values host-key-dsa-private host-key-dsa-public))
    (else (error 'host-key-algorithm->keys "Unsupported host-key-alg ~v" host-key-alg))))

(define (host-key-signature private-key host-key-alg exchange-hash)
  (case host-key-alg
    ((ssh-rsa)
     ;; TODO: offer ssh-rsa. See comment in definition of
     ;; local-algorithm-list in ssh-transport.rkt.
     (error 'host-key-signature "ssh-rsa host key signatures unimplemented"))
    ((ssh-dss)
     (match (asn1-ber-decode-all (pk-sign private-key exchange-hash))
       (`(0 16 ((0 2 ,r-bytes)
		(0 2 ,s-bytes)))
	(bit-string (#"ssh-dss" :: (t:string))
		    ((bit-string ((bs->n r-bytes) :: big-endian integer bits 160)
				 ((bs->n s-bytes) :: big-endian integer bits 160))
		     :: (t:string))))))))

(define (verify-host-key-signature! public-key host-key-alg exchange-hash h-signature)
  ;; TODO: If we are *re*keying, worth checking here that the key hasn't *changed* either.
  (write `(TODO check-host-key!)) (newline) (flush-output) ;; actually check the identity TOFU/POP
  (case host-key-alg
    ((ssh-rsa)
     ;; TODO: offer ssh-rsa. See comment in definition of
     ;; local-algorithm-list in ssh-transport.rkt.
     (error 'verify-host-key-signature! "ssh-rsa host key signatures unimplemented"))
    ((ssh-dss)
     (define signature (bit-string-case h-signature
			 ([ (= #"ssh-dss" :: (t:string #:pack))
			    (r-and-s :: (t:string)) ]
			  (bit-string-case r-and-s
			    ([ (r :: big-endian integer bits 160)
			       (s :: big-endian integer bits 160) ]
			     (asn1-ber-encode `(0 16 ((0 2 ,(n->bs r))
						      (0 2 ,(n->bs s))))))))))
     (when (not (pk-verify public-key exchange-hash signature))
       (error 'verify-host-key-signature! "Signature mismatch")))))

(define (pieces->ssh-host-key pieces)
  (match pieces
    ((struct rsa-public-key (n e))
     (bit-string (#"ssh-rsa" :: (t:string))
		 (e :: (t:mpint))
		 (n :: (t:mpint))))
    ((struct dsa-public-key (y p q g))
     (bit-string (#"ssh-dss" :: (t:string))
		 (p :: (t:mpint))
		 (q :: (t:mpint))
		 (g :: (t:mpint))
		 (y :: (t:mpint))))))

(define (ssh-host-key->pieces blob)
  (bit-string-case blob
    ([ (= #"ssh-rsa" :: (t:string #:pack))
       (e :: (t:mpint))
       (n :: (t:mpint)) ]
     (rsa-public-key n e))
    ([ (= #"ssh-dss" :: (t:string #:pack))
       (p :: (t:mpint))
       (q :: (t:mpint))
       (g :: (t:mpint))
       (y :: (t:mpint)) ]
     (dsa-public-key y p q g))))

;; TODO: proper store for keys

(define (load-private-key filename)
  (pieces->private-key
   (bytes->private-key-pieces
    (base64-decode
     (regexp-replace* #rx"(?m:^-.*-$)"
		      (call-with-input-file filename port->bytes)
		      #"")))))

(define host-key-dsa-private (load-private-key "test-dsa-key"))
(define host-key-dsa-public (pk-key->public-only-key host-key-dsa-private))

(check-equal? (pk-key->datum (pieces->public-key (public-key->pieces host-key-dsa-public))
                             'SubjectPublicKeyInfo)
              (pk-key->datum host-key-dsa-private 'SubjectPublicKeyInfo))