module DNSKey where

import           Data.Binary
import           Data.Binary.Get
import           Data.Binary.Put
import           Data.Bits
import qualified Data.ByteString      as B
import           Data.Int

-- RSA n e
--  n = modulus
--  e = exponent
data RSA = RSA Integer Integer
 deriving (Eq,Show)

getRSA :: Int64 -> Get RSA
getRSA len = do
    elen <- do
        tog <- get
        case tog :: Word8 of
            0 -> get :: Get Word16
            _ -> return $ fromIntegral tog
    ebs <- getByteString (fromIntegral elen)
    nlen <- fmap (fromIntegral . ((-) len)) bytesRead
    nbs <- getByteString nlen
    let e = snd $ B.foldl' (\(c,a) w8 -> (c-1, a + 256^c * fromIntegral w8)) (elen-1,0) ebs
        n = snd $ B.foldl' (\(c,a) w8 -> (c-1, a + 256^c * fromIntegral w8)) (nlen-1,0) nbs
    return $ RSA n e

putRSA :: RSA -> PutM ()
putRSA (RSA n e) = do
    if elen > 255 || elen==0
        then do
            put (0 :: Word8)
            put (fromIntegral elen :: Word16)
        else do
            put (fromIntegral elen :: Word8)
    putByteString ebs
    putByteString nbs
 where
    elen = B.length ebs
    ebs = i2bs_unsized e
    nbs = i2bs_unsized n

-- |@i2bs_unsized i@ converts @i@ to a 'ByteString' of sufficient bytes to express the integer.
-- The integer must be non-negative and a zero will be encoded in one byte.
--
-- Thanks: Thomas DuBuisson (crypto-api)
i2bs_unsized :: Integer -> B.ByteString
i2bs_unsized 0 = B.singleton 0
i2bs_unsized i = B.reverse $ B.unfoldr (\i' -> if i' <= 0 then Nothing else Just (fromIntegral i', (i' `shiftR` 8))) i
{-# INLINE i2bs_unsized #-}

{-
main = do
    bs <- L.getContents
    let rsa = runGet (getRSA (fromIntegral $ L.length bs)) bs
        bs' = runPut (putRSA rsa)
        rsa' = runGet (getRSA (fromIntegral $ L.length bs)) bs'
    print rsa
    print rsa'
-}