;; dig.lisp

(defpackage :dig
  (:use :cl :dns))

(in-package :dig)

(defun dig-class-name (class-id)
  (or (class-id-name class-id)
      class-id))

(defun dig-type-name (type-id)
  (or (type-id-name type-id)
      type-id))

(defgeneric dig-line (record)
  (:method (record)
    (format t "~A~24T~D~32T~A~40T~A~48T~A~%"
            (dig-name record)
            (dig-ttl record) 
            (dig-class record)
            (dig-type record)
            (dig-data record))))

(defgeneric dig-name (record)
  (:method (record)
    (labels-string (name record)))
  (:method ((record question-record))
    (format nil ";~A" (call-next-method))))

(defgeneric dig-ttl (record)
  (:method (record)
    (princ-to-string (ttl record)))
  (:method ((record question-record))
    ""))

(defgeneric dig-type (record)
  (:method (record)
    (dig-type-name (type-id record))))

(defgeneric dig-class (record)
  (:method (record)
    (dig-class-name (class-id record))))

(defgeneric dig-data (record)
  (:method ((record resource-record))
    (princ-to-string (data record)))
  (:method ((record question-record))
    "")
  (:method ((record pointer))
    (labels-string (pointer-name record)))
  (:method ((record name-server))
    (labels-string (name-server record)))
  (:method ((record mail-exchange))
    (format nil "~D ~A" (preference record) (labels-string (exchange record))))
  (:method ((record start-of-authority))
    (format nil "~A ~A ~D ~D ~D ~D ~D"
            (labels-string (main-name record))
            (labels-string (responsible-name record))
            (serial record)
            (refresh record)
            (retry record)
            (expire record)
            (minimum record)))
  (:method ((record address))
    (let ((ip (address record)))
      (format nil "~D.~D.~D.~D"
              (ldb (byte 8 24) ip)
              (ldb (byte 8 16) ip)
              (ldb (byte 8  8) ip)
              (ldb (byte 8  0) ip)))))


(defun dig-message-flags (message)
  (let ((flags '()))
    (macrolet ((add-flag (function symbol)
                 `(when (,function message)
                   (push ',symbol flags))))
      (add-flag recursion-available-p ra)
      (add-flag recursion-desired-p rd)
      (add-flag truncatedp tc)
      (add-flag authoritativep aa)
      (add-flag responsep qr))
    flags))


(defvar *response-codes*
  '((:NOERROR . 0)
    (:FORMAT-ERROR . 1)
    (:SERVFAIL . 2)
    (:NXDOMAIN . 3)
    (:NOTIMPL . 4)
    (:REFUSED . 5)))

(defvar *query-codes*
  '((:QUERY . 0)
    (:IQUERY . 1)
    (:STATUS . 2)))

(defun response-code-name (response-code)
  (or (car (rassoc response-code *response-codes*))
      :UNKNOWN))

(defun query-code-name (query-code)
  (or (car (rassoc query-code *query-codes*))
      :UNKNOWN))

(defmethod dig-display (message)
  (format t ";; ->>HEADER<< opcode: ~A, status: ~A, id: ~A~%"
          (query-code-name (query-code message))
          (response-code-name (response-code message))
          (id message))
  (format t ";; flags: ~(~{~A~^ ~}~); ~
               QUERY: ~D, ANSWER: ~D, AUTHORITY: ~D, ADDITIONAL: ~D~%"
          (dig-message-flags message)
          (question-count message)
          (answer-count message)
          (authority-count message)
          (additional-count message))
  (flet ((dump-section (name entries)
           (when (plusp (length entries))
             (format t "~%;; ~A SECTION:~%" name)
             (loop for entry across entries do
                   (dig-line entry)))))
    (dump-section "QUESTION" (question-records message))
    (dump-section "ANSWER" (answer-records message))
    (dump-section "AUTHORITY" (authority-records message))
    (dump-section "ADDITIONAL" (additional-records message))))

(defun class-id-designator-class-id (designator)
  (etypecase designator
    ((or string symbol) (class-name-id designator))
    (integer designator)))

(defun type-id-designator-type-id (designator)
  (etypecase designator
    ((or string symbol) (type-name-id designator))
    (integer designator)))

(defun make-query (name type class)
  (let ((query (make-instance 'query)))
    (add-question (make-instance 'question-record
                                 :name (string-labels name)
                                 :type-id type
                                 :class-id class)
                  query)
    query))

(defun dig (name server &key (port 53) (type :a)  (class :in))
  (let* ((query (make-query name
                            (type-id-designator-type-id type)
                            (class-id-designator-class-id class)))
         (socket (make-instance 'sb-bsd-sockets:inet-socket
                                :type :datagram
                                :protocol :udp))
         (response-vector (make-array 512 :element-type '(unsigned-byte 8))))
    (flet ((send ()
             (sb-bsd-sockets:socket-send socket (message-vector query) nil
                                         :address (list server port)))
           (receive ()
             (sb-bsd-sockets:socket-receive socket response-vector nil)
             (let ((response (get-message response-vector)))
               (when (query-response-p query response)
                 (dig-display response)
                 t))))
      (block nil
        (dotimes (i 2 (error "No response"))
          (send)
          (ignore-errors
            (sb-ext:with-timeout 2
              (when (receive)
                (return))))))))))