;;; GNU Guix --- Functional package management for GNU
;;; Copyright © 2017 Julien Lepiller <julien@lepiller.eu>
;;;
;;; This file is part of GNU Guix.
;;;
;;; GNU Guix 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.
;;;
;;; GNU Guix 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 GNU Guix.  If not, see <http://www.gnu.org/licenses/>.

(define-module (guix dns)
  #:use-module (ice-9 match)
  #:use-module (ice-9 iconv)
  #:use-module (rnrs bytevectors)
  #:use-module (rnrs arithmetic bitwise)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-9)
  #:export (<dns-flags> make-dns-flags dns-flags?
            dns-flags-response?
            dns-flags-opcode
            dns-flags-authoritative-answer?
            dns-flags-truncation?
            dns-flags-recursion-desired?
            dns-flags-recursion-available?
            dns-flags-rcode

            <dns-query> make-dns-query dns-query?
            dns-query-flags
            dns-query-queries
            dns-query-answers
            dns-query-nameservers
            dns-query-additionals

            <query> make-query query?
            query-name
            query-type
            query-class

            <dns-record> make-dns-record dns-record?
            dns-record-name
            dns-record-type
            dns-record-class
            dns-record-ttl
            dns-record-rdata

            simple-a-query
            dns-query->bytevector
            bytevector->dns-query
            bytevector->ipv4))

;;; Commentary:
;;;
;;; This module provides a DNS implementation. This modules helps construct
;;; valid DNS requests and analyze responses from servers.
;;;
;;; Code:

(define-record-type <dns-flags>
  (make-dns-flags response? opcode authoritative-answer? truncation?
                  recursion-desired? recursion-available? rcode)
  dns-flags?
  (response? dns-flags-response?)
  (opcode dns-flags-opcode)
  (authoritative-answer? dns-flags-authoritative-answer?)
  (truncation? dns-flags-truncation?)
  (recursion-desired? dns-flags-recursion-desired?)
  (recursion-available? dns-flags-recursion-available?)
  (rcode dns-flags-rcode))

(define-record-type <dns-query>
  (make-dns-query flags queries answers nameservers additionals)
  dns-query?
  (flags dns-query-flags)
  (queries dns-query-queries)
  (answers dns-query-answers)
  (nameservers dns-query-nameservers)
  (additionals dns-query-additionals))

(define-record-type <query>
  (make-query name type class)
  query?
  (name query-name)
  (type query-type)
  (class query-class))

(define-record-type <dns-record>
  (make-dns-record name type class ttl rdata)
  dns-record?
  (name dns-record-name)
  (type dns-record-type)
  (class dns-record-class)
  (ttl dns-record-ttl)
  (rdata dns-record-rdata))

(define-record-type <pos-value>
  (make-pos-value pos value)
  pos-value?
  (pos pos-value-pos)
  (value pos-value-value))

;; query type from/to number

(define (type->number type)
  (match type
    ("A" 1)
    ("AAAA" 28)))

(define (type->string type)
  (match type
    (1 "A")
    (28 "AAAA")))

(define (opcode->number opcode)
  (match opcode
    ("QUERY" 0)
    ("IQUERY" 1)
    ("STATUS" 2)))

(define (opcode->string opcode)
  (match opcode
    (0 "QUERY")
    (1 "IQUERY")
    (2 "STATUS")))

(define (rcode->number rcode)
  (match rcode
    ("NOERROR" 0)
    ("FORMATERROR" 1)
    ("SERVFAIL" 2)
    ("NAMEERROR" 3)
    ("NOTIMPLEMENTED" 4)
    ("REFUSED" 5)))

(define (rcode->string rcode)
  (match rcode
    (0 "NOERROR")
    (1 "FORMATERROR")
    (2 "SERVFAIL")
    (3 "NAMEERROR")
    (4 "NOTIMPLEMENTED")
    (5 "REFUSED")))

(define (class->number class)
  (match class
    ("IN" 1)
    ("CS" 2)
    ("CH" 3)
    ("HS" 4)))

(define (class->string class)
  (match class
    (1 "IN")
    (2 "CS")
    (3 "CH")
    (4 "HS")))

(define (write-domain bv components pos)
  "Updates @var{bv} starting at @var{pos} with the @var{components}.
The DNS protocol specifies that each component is preceded by a byte containing
the size of the component, and the last component is followed by the nul byte.
We do not implement the compression algorithm in the query."
  (match components
    ('()
     (begin
       (bytevector-u8-set! bv pos 0)
       (+ pos 1)))
    ((component rest ...)
     (begin
       (bytevector-u8-set! bv pos (string-length component))
       (bytevector-copy! (string->bytevector component "UTF-8") 0
                         bv (+ pos 1) (string-length component))
       (write-domain bv rest (+ pos (string-length component) 1))))))

(define (boolean->number b)
  (if b 1 0))

(define (number->boolean n)
  (not (eq? n 0)))

(define (query-flags->number flags)
  "Returns a number corresponding to the flag bitfield in the DNS header."
  (+ (* 256 128 (boolean->number (dns-flags-response? flags)))
     (* 256 8 (opcode->number (dns-flags-opcode flags)))
     (* 256 4 (boolean->number (dns-flags-authoritative-answer? flags)))
     (* 256 2 (boolean->number (dns-flags-truncation? flags)))
     (* 256   (boolean->number (dns-flags-recursion-desired? flags)))
     (* 128   (boolean->number (dns-flags-recursion-available? flags)))
     (rcode->number (dns-flags-rcode flags))))

(define (create-dns-header flags qdcount ancount nscount arcount)
  "Creates a bytevector containing the header of a DNS query."
  (let ((bv (make-bytevector 12)))
    (bytevector-u16-set! bv 0 15326 (endianness big))
    (bytevector-u16-set! bv 2 (query-flags->number flags) (endianness big))
    (bytevector-u16-set! bv 4 qdcount (endianness big))
    (bytevector-u16-set! bv 6 ancount (endianness big))
    (bytevector-u16-set! bv 8 nscount (endianness big))
    (bytevector-u16-set! bv 10 arcount (endianness big))
    bv))

(define (create-dns-query query)
  "Creates a bytevector containing a question section of a DNS query"
  (let* ((domain (query-name query))
         (len (+ 2 (string-length domain) 4))
         (bv (make-bytevector len)))
    (write-domain bv (string-split domain #\.) 0)
    (bytevector-u16-set! bv (+ 2 (string-length domain))
                         (type->number (query-type query)) (endianness big))
    (bytevector-u16-set! bv (+ 4 (string-length domain))
                         (class->number (query-class query)) (endianness big))
    bv))

(define (create-dns-queries queries)
  (map create-dns-query queries))

;; TODO
(define (create-dns-answers answers)
  '())
(define create-dns-nameservers create-dns-answers)
(define create-dns-additionals create-dns-answers)

(define (dns-query->bytevector query tcp?)
  "Creates a bytevector representing the DNS query to send over the network.
If @code{tcp?} is @code{#t}, the query is suitable for being sent over TCP.
Otherwise, it is suitable to be sent over UDP."
  (let* ((header (create-dns-header
                   (dns-query-flags query)
                   (length (dns-query-queries query))
                   (length (dns-query-answers query))
                   (length (dns-query-nameservers query))
                   (length (dns-query-additionals query))))
         (queries (create-dns-queries (dns-query-queries query)))
         (answers (create-dns-answers (dns-query-answers query)))
         (nameservers (create-dns-answers (dns-query-nameservers query)))
         (additionals (create-dns-answers (dns-query-additionals query)))
         (tcp-header (if tcp? (make-bytevector 2) (make-bytevector 0)))
         (parts-list (append (list tcp-header header) queries answers nameservers additionals))
         (len (fold (lambda (bv l) (+ l (bytevector-length bv))) 0 parts-list))
         (bv (make-bytevector len)))
    (begin
      (if tcp?
        (bytevector-u16-set! tcp-header 0 (- len 2) (endianness big)))
      (fold (lambda (part l)
              (begin
                (bytevector-copy! part 0 bv l (bytevector-length part))
                (+ l (bytevector-length part))))
            0 parts-list)
      bv)))

(define (bytevector->name bv pos)
  "Extracts a name at position @code{pos} in bytevector @code{bv}. This
procedure supports the compression algorithm of DNS names."
  (let* ((component-size (bytevector-u8-ref bv pos))
         (vect (make-bytevector component-size)))
    (if (eq? component-size 0)
        (make-pos-value (+ pos 1) "")
        (begin
          ;; If the first two bytes are 0, the name is not compressed. Otherwise,
          ;; it is compressed and the rest of the field is the position at
          ;; which the complete name can be found.
          (if (eq? (bitwise-and 192 component-size) 0)
              (begin
                (bytevector-copy! bv (+ pos 1)
                                  vect 0 component-size)
                (let ((rest (bytevector->name bv (+ pos 1 component-size))))
                  (make-pos-value (pos-value-pos rest)
                    (string-append (bytevector->string vect "UTF-8") "."
                                 (pos-value-value rest)))))
              (let ((pointer (bitwise-and
                               (bytevector-u16-ref bv pos (endianness big))
                               (- 65535 (* 256 192)))))
                (make-pos-value (+ pos 2)
                  (pos-value-value (bytevector->name bv (+ 2 pointer))))))))))

(define (bytevector->query bv pos)
  (let* ((name (bytevector->name bv pos))
         (type (type->string (bytevector-u16-ref bv (pos-value-pos name)
                                                 (endianness big))))
         (class (class->string (bytevector-u16-ref bv (+ 2 (pos-value-pos name))
                                                   (endianness big)))))
    (make-pos-value (+ 4 (pos-value-pos name))
                    (make-query (pos-value-value name) type class))))

(define (bytevector->queries bv pos num)
  (if (eq? num 0)
    (make-pos-value pos '())
    (let* ((q (bytevector->query bv pos))
           (rest (bytevector->queries bv (pos-value-pos q) (- num 1))))
      (make-pos-value
        (pos-value-pos rest)
        (cons (pos-value-value q)
              (pos-value-value rest))))))

(define (bytevector->dns-records bv pos count)
  (if (> count 0)
      (let* ((result (bytevector->name bv pos))
             (domain (pos-value-value result))
             (npos (pos-value-pos result))
             (type (bytevector-u16-ref bv npos (endianness big)))
             (class (bytevector-u16-ref bv (+ npos 2) (endianness big)))
             (ttl (bytevector-u32-ref bv (+ npos 4) (endianness big)))
             (rdlength (bytevector-u16-ref bv (+ npos 8) (endianness big)))
             (data (make-bytevector rdlength))
             (rest (bytevector->dns-records bv (+ npos 10 rdlength) (- count 1))))
        (bytevector-copy! bv (+ npos 10)
                          data 0 rdlength)
        (make-pos-value (pos-value-pos rest)
          (cons (make-dns-record domain (type->string type)
                                 (class->string class) ttl data)
                (pos-value-value rest))))
      (make-pos-value pos '())))

(define (bytevector->dns-query bv tcp?)
  "Creates a @code{dns-query} object from the @code{bv} bytevector. If @code{tcp?}
is #t, the message is assumed to come from a TCP connection, otherwise it is
treated as if it came from a UDP message."
  (let* ((pos (if tcp? 2 0))
         ;; decode header
         (flags (bytevector-u16-ref bv (+ pos 2) (endianness big)))
         (flags (make-dns-flags
                  (number->boolean (bitwise-and (* 256 128) flags))
                  (opcode->string (/ (bitwise-and (* 256 (+ 8 16 32 64)) flags) (* 256 8)))
                  (number->boolean (bitwise-and (* 256 4) flags))
                  (number->boolean (bitwise-and (* 256 2) flags))
                  (number->boolean (bitwise-and 256 flags))
                  (number->boolean (bitwise-and 128 flags))
                  (rcode->string (bitwise-and 15 flags))))
         (qdcount (bytevector-u16-ref bv (+ pos 4) (endianness big)))
         (ancount (bytevector-u16-ref bv (+ pos 6) (endianness big)))
         (nscount (bytevector-u16-ref bv (+ pos 8) (endianness big)))
         (arcount (bytevector-u16-ref bv (+ pos 10) (endianness big)))
         (pos (+ pos 12))
         (queries (bytevector->queries bv pos qdcount))
         (pos (pos-value-pos queries))
         (answers (bytevector->dns-records bv pos ancount))
         (pos (pos-value-pos answers))
         (nameservers (bytevector->dns-records bv pos nscount))
         (pos (pos-value-pos nameservers))
         (additionals (bytevector->dns-records bv pos arcount)))
    (make-dns-query flags (pos-value-value queries) (pos-value-value answers)
                    (pos-value-value nameservers) (pos-value-value additionals))))

(define (simple-a-query domain)
  "Creates a simple query object that can be passed to @code{dns-query->bytevector}."
  (make-dns-query (make-dns-flags #f "QUERY" #f #f #t #t "NOERROR")
                  (list (make-query domain "A" "IN"))
                  '() '() '()))

(define (bytevector->ipv4 bv)
  "Extracts the rdata section of an A record."
  (string-append
    (number->string (bytevector-u8-ref bv 0)) "."
    (number->string (bytevector-u8-ref bv 1)) "."
    (number->string (bytevector-u8-ref bv 2)) "."
    (number->string (bytevector-u8-ref bv 3))))