initial import.

5 files changed, 414 insertions, 0 deletions

diff --git a/Network/DNS.hs b/Network/DNS.hs
new file mode 100644
index 0000000..223c72d
--- /dev/null
+++ b/Network/DNS.hs
@@ -0,0 +1,4 @@
+module Network.DNS (parseResponse, composeQuery) where
+
+import Network.DNS.Query
+import Network.DNS.Response
diff --git a/Network/DNS/Query.hs b/Network/DNS/Query.hs
new file mode 100644
index 0000000..8fcf694
--- /dev/null
+++ b/Network/DNS/Query.hs
@@ -0,0 +1,78 @@
+module Network.DNS.Query (composeQuery) where
+
+import Data.ByteString.Lazy (ByteString)
+import Data.Char
+import Network.DNS.StateBinary
+import Network.DNS.Types
+
+----------------------------------------------------------------
+
+composeQuery :: Int -> [Question] -> ByteString
+composeQuery idt qs = runSPut (encodeQuery qry)
+  where
+    hdr = header defaultQuery
+    qry = defaultQuery {
+        header = hdr {
+           identifier = idt
+         , qdCount = length qs
+         }
+      , question = qs
+      }
+
+----------------------------------------------------------------
+
+encodeQuery :: DNSFormat -> SPut
+encodeQuery fmt = do
+  let hdr = header fmt
+      qs = question fmt
+  encodeHeader hdr
+  encodeQuestion qs
+  return ()
+
+encodeHeader :: DNSHeader -> SPut
+encodeHeader hdr = do
+    encodeIdentifier $ identifier hdr
+    encodeFlags $ flags hdr
+    decodeQdCount $ qdCount hdr
+    decodeAnCount $ anCount hdr
+    decodeNsCount $ nsCount hdr
+    decodeArCount $ arCount hdr
+  where
+    encodeIdentifier = putInt16
+    decodeQdCount = putInt16
+    decodeAnCount = putInt16
+    decodeNsCount = putInt16
+    decodeArCount = putInt16
+
+encodeFlags :: DNSFlags -> SPut
+encodeFlags _ = put16 0x0100 -- xxx
+
+encodeQuestion :: [Question] -> SPut
+encodeQuestion qs = do
+  let q = head qs 
+      dom = qname q
+      typ = qtype q
+  encodeDomain dom
+  putInt16 . typeToInt $ typ
+  put16 1
+
+----------------------------------------------------------------
+
+encodeDomain :: Domain -> SPut
+encodeDomain dom = do
+    let ss = split '.' dom
+        ls = map length ss
+    mapM_ encodeSubDomain $ zip ls ss
+    put8 0
+  where
+    encodeSubDomain (len,sub) = do
+      putInt8 len
+      mapM_ (putInt8 . ord) sub
+
+split :: Char -> String -> [String]
+split _ "" = []
+split c cs
+  | null rest = s : split c rest
+  | otherwise = s : split c (tail rest)
+  where
+    (s,rest) = break (c ==) cs
diff --git a/Network/DNS/Response.hs b/Network/DNS/Response.hs
new file mode 100644
index 0000000..c6e0410
--- /dev/null
+++ b/Network/DNS/Response.hs
@@ -0,0 +1,133 @@
+module Network.DNS.Response (parseResponse) where
+
+import Control.Monad
+import Data.Bits
+import Data.ByteString.Lazy (ByteString)
+import Data.Char
+import Data.IP
+import Network.DNS.StateBinary
+import Network.DNS.Types
+
+----------------------------------------------------------------
+
+parseResponse :: ByteString -> DNSFormat
+parseResponse bs = runSGet decodeResponse bs
+
+----------------------------------------------------------------
+
+decodeResponse :: SGet DNSFormat
+decodeResponse = do
+    hd <- decodeHeader
+    DNSFormat hd <$> (decodeQueries $ qdCount hd)
+                 <*> (decodeRRs $ anCount hd)
+                 <*> (decodeRRs $ nsCount hd)
+                 <*> (decodeRRs $ arCount hd)
+
+----------------------------------------------------------------
+
+decodeFlags :: SGet DNSFlags
+decodeFlags = do
+  flgs <- get16
+  return $ DNSFlags (getQorR flgs)
+                    (getOpcode flgs)
+                    (getAuthAnswer flgs)
+                    (getTrunCation flgs)
+                    (getRecDesired flgs)
+                    (getRecAvailable flgs)
+                    (getRcode flgs)
+  where
+    getQorR w = if testBit w 15 then QR_Response else QR_Query
+    getOpcode w = toEnum $ fromIntegral $ shiftR w 11 .&. 0x0f
+    getAuthAnswer w = testBit w 10
+    getTrunCation w = testBit w 9
+    getRecDesired w = testBit w 8
+    getRecAvailable w = testBit w 7
+    getRcode w = toEnum $ fromIntegral $ w .&. 0x0f
+
+----------------------------------------------------------------
+
+decodeHeader :: SGet DNSHeader
+decodeHeader = DNSHeader <$> decodeIdentifier
+                         <*> decodeFlags
+                         <*> decodeQdCount
+                         <*> decodeAnCount
+                         <*> decodeNsCount
+                         <*> decodeArCount
+  where
+    decodeIdentifier = getInt16
+    decodeQdCount = getInt16
+    decodeAnCount = getInt16
+    decodeNsCount = getInt16
+    decodeArCount = getInt16
+
+----------------------------------------------------------------
+
+decodeQueries :: Int -> SGet [Question]
+decodeQueries n = replicateM n decodeQuery
+
+decodeType :: SGet TYPE
+decodeType = intToType <$> getInt16
+
+decodeQuery :: SGet Question
+decodeQuery = Question <$> decodeDomain
+                       <*> (decodeType <* ignoreClass)
+
+decodeRRs :: Int -> SGet [ResourceRecord]
+decodeRRs n = replicateM n decodeRR
+
+decodeRR :: SGet ResourceRecord
+decodeRR = do
+    Question dom typ <- decodeQuery
+    ttl <- decodeTTL
+    len <- decodeRLen
+    dat <- decodeRData typ len
+    return ResourceRecord { rrname = dom
+                          , rrtype = typ
+                          , rrttl  = ttl
+                          , rdlen  = len
+                          , rdata  = dat
+                          }
+  where
+    decodeTTL = fromIntegral <$> get32
+    decodeRLen = getInt16
+
+decodeRData :: TYPE -> Int -> SGet RDATA
+decodeRData NS _   = RD_NS <$> decodeDomain
+decodeRData A len  = (RD_A . toIPv4) <$> getNBytes len
+decodeRData AAAA len  = (RD_AAAA . toIPv6 . combine) <$> getNBytes len
+  where
+    combine [] = []
+    combine [_] = error "combine"
+    combine (a:b:cs) =  a * 256 + b : combine cs
+decodeRData _  len = RD_OTH <$> getNBytes len
+
+----------------------------------------------------------------
+
+decodeDomain :: SGet Domain
+decodeDomain = do
+    pos <- getPosition
+    c <- getInt8
+    if c == 0
+      then return ""
+      else do
+        let n = getValue c
+        if isPointer c
+          then do
+            d <- getInt8
+            let offset = n * 256 + d
+            maybe (error "decodeDomain") id <$> pop offset
+          else do
+            hs <- decodeString n
+            ds <- decodeDomain
+            let dom = hs ++ "." ++ ds
+            push pos dom
+            return dom
+  where
+    getValue c = c .&. 0x3f
+    isPointer c = testBit c 7 && testBit c 6
+
+decodeString :: Int -> SGet String
+decodeString n = map chr <$> getNBytes n
+
+ignoreClass :: SGet ()
+ignoreClass = () <$ get16
diff --git a/Network/DNS/StateBinary.hs b/Network/DNS/StateBinary.hs
new file mode 100644
index 0000000..3b89665
--- /dev/null
+++ b/Network/DNS/StateBinary.hs
@@ -0,0 +1,107 @@
+module Network.DNS.StateBinary where
+
+import Control.Monad.State
+import Data.Binary.Get
+import Data.Binary.Put
+import Data.ByteString.Lazy (ByteString)
+import qualified Data.ByteString.Lazy as BS hiding (ByteString)
+import Data.IntMap (IntMap)
+import qualified Data.IntMap as IM (insert, lookup, empty)
+import Data.Word
+import Network.DNS.Types
+import Prelude hiding (lookup)
+
+----------------------------------------------------------------
+  
+type SGet = StateT PState Get
+
+type PState = IntMap Domain
+
+----------------------------------------------------------------
+
+(<$>) :: (Monad m) => (a -> b) -> m a -> m b
+(<$>) = liftM
+
+(<$) :: (Monad m) => b -> m a -> m b
+x <$ y = y >> return x
+
+(<*>) :: (Monad m) => m (a -> b) -> m a -> m b
+(<*>) = ap
+
+(<*) :: (Monad m) => m a -> m b -> m a
+(<*) ma mb = do
+    a <- ma
+    mb
+    return a
+
+----------------------------------------------------------------
+
+type SPut = Put
+
+put8 :: Word8 -> SPut
+put8  = putWord8
+
+put16 :: Word16 -> SPut
+put16 = putWord16be
+
+put32 :: Word32 -> SPut
+put32 = putWord32be
+
+putInt8 :: Int -> SPut
+putInt8  = put8  . fromIntegral
+
+putInt16 :: Int -> SPut
+putInt16 = put16 . fromIntegral
+
+putInt32 :: Int -> SPut
+putInt32 = put32 . fromIntegral
+
+----------------------------------------------------------------
+
+get8 :: SGet Word8
+get8  = lift getWord8
+
+get16 :: SGet Word16
+get16 = lift getWord16be
+
+get32 :: SGet Word32
+get32 = lift getWord32be
+
+getInt8 :: SGet Int
+getInt8 = fromIntegral <$> get8
+
+getInt16 :: SGet Int
+getInt16 = fromIntegral <$> get16
+
+getInt32 :: SGet Int
+getInt32 = fromIntegral <$> get32
+
+----------------------------------------------------------------
+
+getPosition :: SGet Int
+getPosition = fromIntegral <$> lift bytesRead
+
+getNBytes :: Int -> SGet [Int]
+getNBytes len = toInts <$> getNbytes len
+  where
+    toInts = map fromIntegral . BS.unpack
+    getNbytes = lift . getLazyByteString . fromIntegral
+
+----------------------------------------------------------------
+
+push :: Int -> Domain -> SGet ()
+push n d = modify (\m -> IM.insert n d m)
+
+pop :: Int -> SGet (Maybe Domain)
+pop n = IM.lookup n <$> get
+
+----------------------------------------------------------------
+
+initialState :: IntMap Domain
+initialState = IM.empty
+
+runSGet :: SGet DNSFormat -> ByteString -> DNSFormat
+runSGet res bs = fst $ runGet (runStateT res initialState) bs
+
+runSPut :: Put -> ByteString
+runSPut = runPut
\ No newline at end of file
diff --git a/Network/DNS/Types.hs b/Network/DNS/Types.hs
new file mode 100644
index 0000000..aa32431
--- /dev/null
+++ b/Network/DNS/Types.hs
@@ -0,0 +1,92 @@
+module Network.DNS.Types where
+
+import Data.Maybe
+import Data.IP
+
+data TYPE = A | AAAA | NS | TXT | MX | UNKNOWN deriving (Eq, Show)
+
+rrDB :: [(TYPE, Int)]
+rrDB = [
+    (A,     1)
+  , (NS,    2)
+  , (MX,   15)
+  , (TXT,  16)
+  , (AAAA, 28)
+  ]
+
+rookup                  :: (Eq b) => b -> [(a,b)] -> Maybe a
+rookup _    []          =  Nothing
+rookup  key ((x,y):xys)
+  | key == y          =  Just x
+  | otherwise         =  rookup key xys
+
+intToType :: Int -> TYPE
+intToType n = maybe UNKNOWN id $ rookup n rrDB
+typeToInt :: TYPE -> Int
+typeToInt t = maybe 0 id $ lookup t rrDB
+
+data QorR = QR_Query | QR_Response deriving (Eq, Show)
+
+data OPCODE = OP_STD | OP_INV | OP_SSR deriving (Eq, Show, Enum)
+
+data RCODE = NoErr | FormatErr | ServFail | NameErr | NotImpl | Refused deriving (Eq, Show, Enum)
+
+type Domain = String
+
+data Question = Question {
+    qname  :: Domain
+  , qtype  :: TYPE
+  } deriving (Eq, Show)
+
+data ResourceRecord = ResourceRecord {
+    rrname :: Domain
+  , rrtype :: TYPE
+  , rrttl  :: Int
+  , rdlen  :: Int
+  , rdata  :: RDATA
+  } deriving (Eq, Show)
+
+data RDATA = RD_NS Domain | RD_A IPv4 | RD_AAAA IPv6 | RD_OTH [Int] deriving (Eq, Show)
+
+data DNSFlags = DNSFlags {
+    qOrR         :: QorR
+  , opcode       :: OPCODE
+  , authAnswer   :: Bool
+  , trunCation   :: Bool
+  , recDesired   :: Bool
+  , recAvailable :: Bool
+  , rcode        :: RCODE
+  } deriving (Eq, Show)
+
+data DNSHeader = DNSHeader {
+    identifier :: Int
+  , flags      :: DNSFlags
+  , qdCount    :: Int
+  , anCount    :: Int
+  , nsCount    :: Int
+  , arCount    :: Int
+  } deriving (Eq, Show)
+
+data DNSFormat = DNSFormat {
+    header     :: DNSHeader
+  , question   :: [Question]
+  , answer     :: [ResourceRecord]
+  , authority  :: [ResourceRecord]
+  , additional :: [ResourceRecord]
+  } deriving (Eq, Show)
+
+defaultQuery :: DNSFormat
+defaultQuery = DNSFormat {
+    header = DNSHeader {
+       identifier = 0
+     , flags = undefined
+     , qdCount = 0
+     , anCount = 0
+     , nsCount = 0
+     , arCount = 0
+     }
+  , question   = []
+  , answer     = []
+  , authority  = []
+  , additional = []
+  }