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

(require bitsyntax)
(require racket/unix-socket)
(require net/base64)

(define-values (i o) (unix-socket-connect (getenv "SSH_AUTH_SOCK")))

(define SSH_AGENT_FAILURE 5)
(define SSH_AGENT_SUCCESS 6)
(define SSH_AGENTC_REQUEST_IDENTITIES 11)
(define SSH_AGENT_IDENTITIES_ANSWER 12)
(define SSH_AGENTC_SIGN_REQUEST 13)
(define SSH_AGENT_SIGN_RESPONSE 14)
(define SSH_AGENTC_ADD_IDENTITY 17)
(define SSH_AGENTC_REMOVE_IDENTITY 18)
(define SSH_AGENTC_REMOVE_ALL_IDENTITIES 19)
(define SSH_AGENTC_ADD_SMARTCARD_KEY 20)
(define SSH_AGENTC_REMOVE_SMARTCARD_KEY 21)
(define SSH_AGENTC_LOCK 22)
(define SSH_AGENTC_UNLOCK 23)
(define SSH_AGENTC_ADD_ID_CONSTRAINED 25)
(define SSH_AGENTC_ADD_SMARTCARD_KEY_CONSTRAINED 26)
(define SSH_AGENTC_EXTENSION 27)
(define SSH_AGENT_EXTENSION_FAILURE 28)

(struct identity (blob comment) #:transparent)

(define (write-packet o type bs)
  (write-bytes (bit-string->bytes
                (bit-string ((+ 1 (bytes-length bs)) :: bits 32)
                            (type :: bits 8)
                            (bs :: binary)))
               o)
  (flush-output o))

(define (read-packet i)
  (bit-string-case (read-bytes 4 i)
    ([(len :: bits 32)]
     (bit-string-case (read-bytes len i)
       ([(type :: bits 8) (body :: binary)]
        (values type body))))))

(define (list-keys i o)
  (write-packet o SSH_AGENTC_REQUEST_IDENTITIES #"")
  (define-values (response-type body) (read-packet i))
  (when (not (= response-type SSH_AGENT_IDENTITIES_ANSWER))
    (error 'list-keys "Invalid response from SSH agent: ~a" response-type))
  (bit-string-case body
    ([ (nkeys :: bits 32) (body :: binary) ]
     (let loop ((acc-rev '()) (nkeys nkeys) (body body))
       (if (zero? nkeys)
           (reverse acc-rev)
           (bit-string-case body
             ([ (bloblen :: bits 32) (blob :: binary bytes bloblen)
                (commentlen :: bits 32) (comment :: binary bytes commentlen)
                (rest :: binary) ]
              (loop (cons (identity (bit-string->bytes blob)
                                    (bytes->string/utf-8 (bit-string->bytes comment)))
                          acc-rev)
                    (- nkeys 1)
                    rest))))))))

(define (blob-ed25519-key blob)
  (bit-string-case blob
    ([ (= 11 :: bits 32) (= #"ssh-ed25519" :: binary bytes 11)
       (= 32 :: bits 32) (pk :: binary bytes 32) ]
     (bit-string->bytes pk))
    (else #f)))

(define (sign data id i o)
  (write-packet o SSH_AGENTC_SIGN_REQUEST
                (bit-string->bytes
                 (bit-string ((bytes-length (identity-blob id)) :: bits 32)
                             ((identity-blob id) :: binary)
                             ((bytes-length data) :: bits 32)
                             (data :: binary)
                             (0 :: bits 32))))
  (define-values (response-type body) (read-packet i))
  (when (not (= response-type SSH_AGENT_SIGN_RESPONSE))
    (error 'sign "Invalid response from SSH agent: ~a" response-type))
  (bit-string-case body
    ([ (len :: bits 32) (signature :: binary bytes len) ]
     (bit-string->bytes signature))))

(let ((ids (filter (lambda (i) (blob-ed25519-key (identity-blob i))) (list-keys i o))))
  (for-each writeln ids)
  (newline)
  )