# dns query using Udp extension
# Colin McCormack

lappend auto_path [pwd]/Udp
package require Udp
package require TclOO
package provide Dns 1.1

oo::class create DNS {
    method decodeName {varname} {
        upvar 1 $varname text
        set line {}
        set suffix ""
        while {$text ne ""} {
            binary scan $text ca* len text
            if {$len < 0} {
                # compression
                binary scan $text ca* offset text
                set offset [expr {($len & 0x3F)<<8 | ($offset & 0xFF)}]
                variable response
                set reference [string range $response $offset end]
                set suffix .[my decodeName reference]
                break
            } elseif {$len} {
                lappend line [string range $text 0 $len-1]
                set text [string range $text $len end]
            } else {
                break
            }
        }
        return [join $line .]$suffix
    }

    method decodeQD {varname qdcount} {
        upvar 1 $varname text
        set result {}
        while {$qdcount} {
            incr qdcount -1
            set name [my decodeName text]
            binary scan $text SSa* qtype qclass text
            lappend result [list $name $qtype $qclass]
        }
        return $result
    }

    method decodeAddress {address} {
        catch {
            binary scan $address cccc a1 a2 a3 a4
            foreach v {a1 a2 a3 a4} {
                set $v [expr { [set $v] & 0xff }]
            }
            set result [join [list $a1 $a2 $a3 $a4] .]
        } e eo
        return $result
    }

    method rdA {rd} {
        # a host address
        return [my decodeAddress $rd]
    }

    method rdNS {rd} {
        # an authoritative name server
        return [my decodeName rd]
    }

    method rdMD {rd} {
        # a mail destination (Obsolete - use MX)
        return [my decodeName rd]
    }

    method rdMF {rd} {
        # a mail forwarder (Obsolete - use MX)
        return [my decodeName rd]
    }

    method rdCNAME {rd} {
        # the canonical name for an alias
        return [my decodeName rd]
    }

    method rdSOA {rd} {
        # marks the start of a zone of authority
        set mname [my decodeName rd]
        set rname [my decodeName rd]
        binary scan $rd IIIII serial refresh retry expire minimum
        return [list mname $mname rname $rname serial $serial refresh $refresh retry $retry expire $expire minimum $minimum]
    }

    method rdMB {rd} {
        # a mailbox domain name (EXPERIMENTAL)
        return [my decodeName rd]
    }

    method rdMG {rd} {
        # a mail group member (EXPERIMENTAL)
        return [my decodeName rd]
    }

    method rdMR {rd} {
        # a mail rename domain name (EXPERIMENTAL)
        return [my decodeName rd]
    }

    method rdNULL {rd} {
        # a null RR (EXPERIMENTAL)
        return $rd
    }

    method rdWKS {rd} {
        # a well known service description
        binary scan $rd Ica* address protocol rd
        set address [my decodeAddress $rd]
        return [list address $address protocol $protocol bitmap $rd]
    }

    method rdPTR {rd} {
        # a domain name pointer
        return [my decodeName rd]
    }

    method rdHINFO {rd} {
        # host information
        set cpu [my decodeName rd]
        set os [my decodeName rd]
        return [list cpu $cpu os $os]
    }

    method rdMINFO {rd} {
        # mailbox or mail list information
        set rmailbx [my decodeName rd]
        set emailbx [my decodeName rd]
        return [list rmailbx $rmailbx emailbx $emailbx]
    }

    method rdMX {rd} {
        # mail exchange
        binary scan $rd S preference rd
        return [list preference $preference [my decodeName rd]]
    }

    method rdTXT {rd} {
        # text strings
        set result {}
        while {$rd ne ""} {
            lappend result [my decodeName rd]
        }
        return $result
    }

    method rr {varname count} {
        upvar 1 $varname text
        set result {}

        while {$count && $text ne ""} {
            incr count -1
            set name [my decodeName text]
            binary scan $text SSISa* type class ttl rdlength text
            set type [string trimleft $type 0]
            variable qtypes; catch {set type [dict get $qtypes $type]}
            set class [string trimleft $class 0]
            variable qclasses; catch {set class [dict get $qclasses $class]}

            set rdata [string range $text 0 $rdlength-1]
            catch {set rdata [my rd$type $rdata]}

            set text [string range $text $rdlength end]

            lappend result [list name $name type $type class $class ttl $ttl rdata $rdata]
        }

        return $result
    }

    method decode {payload from port} {
        variable response $payload

        binary scan $payload SB8B8SSSS id h1 h2 qdcount ancount nscount arcount
        lassign [split $h1 ""] qr . . . . aa tc rd
        set rcode 0b0[string trimleft [string range $h2 4 end] 0]
        set rcode [expr {$rcode + 0}]
        set ra [string index $h2 0]
        set id [string trimleft $id 0]
        set result [list id $id qr $qr aa $aa tc $tc rd $rd ra $ra rcode $rcode qdcount $qdcount ancount $ancount nscount $nscount arcount $arcount]
        if {$rcode} {
            variable rcodes
            dict set result error [dict get $rcodes $rcode]
        }

        set content [string range $payload 12 end]

        dict set result qd [my decodeQD content $qdcount]
        dict set result an [my rr content $ancount]
        dict set result ns [my rr content $nscount]
        dict set result ar [my rr content $arcount]
        if {$content ne ""} {
            dict set result remainder $content
        }
        #puts stderr "RESPONSE: $result"
        return $result
    }

    # response - get result and invoke callback
    method response {payload from port chan} {
        if {[catch {
            my decode $payload $from $port
        } result eo]} {
            puts stderr "$result ($eo)"
        } else {
            # got a complete response - invoke callback for it.
            variable callbacks
            set id [dict get $result id]
            if {[dict exists $callbacks $id]} {
                set callback [dict get $callbacks $id]
                dict unset callbacks $id
                {*}$callback $result
            }
        }
    }

    method qhead {id {opcode QUERY} {recurse 1}} {
        if {$recurse} {
            set h 0x0100
        } else {
            set h 0
        }
        switch $opcode {
            QUERY {}
            IQUERY {
                set h [expr {$h | 0x10}]
            }
            STATUS {
                set h [expr {$h | 0x100}]
            }
            default {
                error "opcode must be one of QUERY, IQUERY, STATUS"
            }
        }
        return [binary format SS $id $h]
    }

    method dquery {name {qtype A} {qclass IN}} {
        variable qtypes; variable qclasses
        set query ""
        foreach label [split $name .] {
            set len [string length $label]
            if {$len > 255} {
                error "name component too long: '$label'"
            }
            append query [binary format ca$len $len $label]
        }
        append query \0
        append query [binary format SS [dict get $qtypes $qtype] [dict get $qclasses $qclass]]
        return $query
    }

    method dqueries {id args} {
        set count 0
        set opcode QUERY
        set qtype A
        set qclass *
        set recurse 1
        set query ""
        set skip 0

        foreach arg $args {
            if {$skip} {
                set $var $arg
                set skip 0
                continue
            }
            switch $arg {
                -qtype -
                -qclass -
                -opcode -
                -callback -
                -recurse {
                    incr skip
                    set var [string trim $arg -]
                }
                default {
                    append query [my dquery $arg $qtype $qclass]
                    incr count
                }
            }
        }

        if {[info exists callback]} {
            variable callbacks; dict set callbacks $id $callback
        }

        set query "[my qhead $id $opcode $recurse][binary format SSSS $count 0 0 0]$query"
        return $query
    }

    method query {dns args} {
        variable udp; variable domain
        variable qcount; incr qcount
        set query [my dqueries $qcount {*}$args]
        udp::send $udp $dns $domain $query
    }

    destructor {
        variable udp
        catch {chan close $udp}
    }

    constructor {args} {
        variable domain 53
        variable {*}$args

        variable qcount 1
        variable udp [::udp create 0 [list [self] response]]

        variable callbacks {}
        variable rcodes {
            0 {No error condition}
            1 {Format error - The name server was unable to interpret the query.}
            2 {Server failure - The name server was unable to process this query due to a problem with the name server.}
            3 {Name Error - Meaningful only for responses from an authoritative name server, this code signifies that the domain name referenced in the query does not exist.}
            4 {Not Implemented - The name server does not support the requested kind of query.}
            5 {Refused - The name server refuses to perform the specified operation for policy reasons.  For example, a name server may not wish to provide the information to the particular requester, r a name server may not wish to perform a particular operation (e.g., zone transfer) for particular data.}
        }

        variable qtypes
        foreach {n v d} {
            A 1 {a host address}
            NS 2 {an authoritative name server}
            MD 3 {a mail destination (Obsolete - use MX)}
            MF 4 {a mail forwarder (Obsolete - use MX)}
            CNAME 5 {the canonical name for an alias}
            SOA 6 {marks the start of a zone of authority}
            MB 7 {a mailbox domain name (EXPERIMENTAL)}
            MG 8 {a mail group member (EXPERIMENTAL)}
            MR 9 {a mail rename domain name (EXPERIMENTAL)}
            NULL 10 {a null RR (EXPERIMENTAL)}
            WKS 11 {a well known service description}
            PTR 12 {a domain name pointer}
            HINFO 13 {host information}
            MINFO 14 {mailbox or mail list information}
            MX 15 {mail exchange}
            TXT 16 {text strings}
            AXFR 252 {A request for a transfer of an entire zone}
            MAILB 253 {A request for mailbox-related records (MB, MG or MR)}
            MAILA 254 {A request for mail agent RRs (Obsolete - see MX)}
            * 255 {A request for all records}
        } {
            dict set qtypes $n $v
            dict set qtypes $v $n
        }

        variable qclasses 
        foreach {n v d} {
            IN 1 {the Internet}
            CS 2 {the CSNET class (Obsolete - used only for examples in some obsolete RFCs)}
            CH 3 {the CHAOS class}
            HS 4 {Hesiod [Dyer 87]}
            *               255 {any class}
        } {
            dict set qclasses $n $v
            dict set qclasses $v $n
        }
    }
}

oo::class create DNSsimple {
    superclass DNS

    method response {payload from port} {
        if {[catch {
            my decode $payload $from $port
        } result eo]} {
            puts stderr "$result ($eo)"
        } else {
            # got a complete response - invoke callback for it.
            variable callbacks
            set id [dict get $result id]
            if {[dict exists $callbacks $id]} {
                if {[dict exists $result error]} {
                    # got an error
                    set lookup [list [dict get $result error]]
                } else {
                    # post-process the result to just give simplified answers
                    set lookup {}
                    foreach {v} [dict get $result an] {
                        dict lappend lookup [string trimleft [dict get $v name] .] [dict get $v rdata]
                    }
                }

                set callback [dict get $callbacks $id]
                dict unset callbacks $id
                {*}$callback {*}$lookup
            }
        }
    }

    method reverse {dns args} {
        set skip 0
        set rquery {-qtype PTR}
        foreach arg $args {
            if {$skip} {
                set skip 0
            } elseif {[string match -* $arg]} {
                if {$arg eq "-qtype"} {
                    error "reverse queries cannot specify -qtype"
                }
                set skip 1
            } else {
                # got to be an IP address
                set arg [join [lreverse [split $arg .]] .].in-addr.arpa
            }
            lappend rquery $arg
        }
        my query $dns {*}$rquery
    }

    constructor {args} {
        next {*}$args
    }
}

if {[info exists argv0] && ($argv0 eq [info script])} {
    set dns 192.168.178.2        ;# your dns server

    proc putss {where args} {
        puts $where DNS:$args
    }

    # simple DNS query processing
    DNSsimple create dnss
    dnss query $dns -callback {putss stderr} google.com        ;# make a simple query
    dnss reverse $dns -callback {putss stderr} 74.125.237.48
    dnss query $dns -callback {putss stderr} localhost
    dnss query $dns -callback {putss stderr} google.com        localhost

    # full DNS query processing
    DNS create dns
    dns query $dns -callback {puts stderr} thighbone        ;# make a query
    dns query $dns -qtype A -callback {puts stderr} google.com        ;# make a query

    vwait forever
}