My simple module that I use for send/recv xor-encrypted data via sockets. #haskell #socket #xor

SockEncrypted.hs

{-# LANGUAGE OverloadedStrings #-}
module SockEncrypted
     ( SockEnc (..)
     , nc2
     , chunksof
     , conn2
     , recvconn
     , SockAddr
     , EncSock
     ) where

import           Control.Monad                  (void)
import           Data.Bits                      (xor)
import qualified Data.ByteString.Lazy.Char8     as L8 (ByteString, break,
                                                       concat, cycle, drop,
                                                       init, length, lines,
                                                       pack, readInteger, tail,
                                                       take, unlines, zipWith)
import           Data.Char                      (chr, ord)
import           Data.Int                       (Int64)
import           Data.Maybe                     (fromJust)
import           Data.Monoid                    ((<>))
import           Data.UnixTime                  (getUnixTime, utMicroSeconds,
                                                 utSeconds)
import           Network.Socket                 (Family (AF_INET), PortNumber,
                                                 SockAddr (SockAddrInet),
                                                 Socket,
                                                 SocketOption (ReuseAddr),
                                                 SocketType (Stream), accept,
                                                 addrAddress, addrFamily, bind,
                                                 close, connect,
                                                 defaultProtocol, getAddrInfo,
                                                 iNADDR_ANY, listen,
                                                 setSocketOption, socket,
                                                 withSocketsDo)
import qualified Network.Socket.ByteString.Lazy as NL (recv, send)

-- tests
-- import qualified Data.ByteString.Lazy.Char8     as L8 (hPutStr, putStr)
-- import           System.IO
--
-- fix f = let x = f x in x
--
-- filerecvtest = recvconn "pizdec228" 5003 8192 >>= \(cn,addr) -> do
--   print $ "connected from: " <> show addr
--   withBinaryFile "afile" WriteMode $ \h -> do
--     recvparts cn (L8.hPutStr h)
--     seclose cn
--
-- echoservertest = recvconn "pizdec228" 5003 8192 >>= \(cn,addr) -> do
--   print $ "connected from: " <> show addr
--   fix $ \e -> recv cn >>= send cn >> e
--
-- echoclienttest = do
--   L8.putStr "enter password:  " -- "pizdec228"
--   pw <- fmap L8.pack getLine
--   putStrLn "  - Ok. Trying to connect..."
--   mbcn <- conn2 "localhost" "5003" pw
--   case mbcn of
--     Nothing -> putStrLn "incorrect password!" >> echoclienttest
--     Just cn -> do
--       putStrLn "connected!"
--       fix $ \e -> do
--         x <- fmap L8.pack getLine
--         send cn x >>= print
--         recv cn   >>= print
--         e
-- tests

class SockEnc a where
  sk            :: a -> Socket
  nckey         :: a -> L8.ByteString
  chunksize     :: a -> Int64
  bnd           :: a -> L8.ByteString

  send          :: a -> L8.ByteString -> IO Int64
  send senc msg = sendloop mchunks where
    k = nckey senc
    s = sk senc
    mchunks = chunksof (chunksize senc - 512) msg
    sendloop [] = sendx k s ["-"] >> return 0
    sendloop (c:cs) = do
      x  <- sendx k s (bnd senc : [c])
      _  <- recvx k s 8
      xs <- sendloop cs
      return $ x + xs - L8.length (bnd senc) - 1

  recv          :: a -> IO L8.ByteString
  recv = fmap L8.concat .recvloop where
    recvloop senc = recv0 senc >>= maybe
      (return [])
      (\part0 -> do
        part1 <- recvloop senc
        return (part0 : part1))

  recvparts     :: a -> (L8.ByteString -> IO ()) -> IO ()
  recvparts senc func = recv0 senc >>= maybe
    (return ())
    (\part0 -> func part0 >> recvparts senc func)

  seclose       :: a -> IO ()
  seclose = close .sk


recv0 :: SockEnc a => a -> IO (Maybe L8.ByteString)
recv0 senc = do
  (mbbound,part) <- L8.break (=='\n')
                <$> recvx (nckey senc) (sk senc) (chunksize senc)
  if mbbound == bnd senc then do
    _ <- sendx (nckey senc) (sk senc) ["-"]
    return .Just $ L8.tail part
  else return Nothing

boundary :: IO L8.ByteString
boundary = getUnixTime >>= \ut ->
  return $ (L8.pack.show $ utSeconds ut) <>
           (L8.pack.show $ utMicroSeconds ut)

chunksof :: Int64 -> L8.ByteString -> [L8.ByteString]
chunksof i bs
  | bs == mempty = []
  | otherwise = L8.take i bs : chunksof i (L8.drop i bs)

sendx :: L8.ByteString -> Socket -> [L8.ByteString] -> IO Int64
sendx k sok bss = NL.send sok $ nc2 k line where
  line = if null bss then mempty else L8.init (L8.unlines bss)

recvx :: L8.ByteString -> Socket -> Int64 -> IO L8.ByteString
recvx k s = fmap (nc2 k) .NL.recv s

nc2 :: L8.ByteString -> L8.ByteString -> L8.ByteString
nc2 key x
  | key == mempty = x
  | otherwise = L8.pack $ L8.zipWith enc (L8.cycle key) x
  where enc c1 c2 = chr $ ord c1 `xor` ord c2

conn2 :: String -> String -> L8.ByteString -> IO (Maybe EncSock)
conn2 h p k = withSocketsDo $ do
  addr <- head <$> getAddrInfo Nothing (Just h) (Just p)
  sock <- socket (addrFamily addr) Stream defaultProtocol
  connect sock (addrAddress addr)
  void $ sendx "1337" sock [k]
  answer <- L8.lines <$> recvx k sock 512
  case answer of
    ["on",blsz,bound] -> do
      let sz = fromJust $ L8.readInteger blsz
      void $ sendx k sock ["on"]
      return .Just $ EncSock sock k (fromInteger $ fst sz) bound
    _ -> close sock >> return Nothing

recvconn :: L8.ByteString -> PortNumber -> Int64 -> IO (EncSock,SockAddr)
recvconn ki p sz = withSocketsDo $ do
  asock <- socket AF_INET Stream 0
  setSocketOption asock ReuseAddr 1
  bind asock (SockAddrInet p iNADDR_ANY)
  listen asock 5
  wait4 asock where
    wait4 asock = do
      (sock,a) <- accept asock
      k <- recvx "1337" sock 512
      if k == ki then do
        bound <- boundary
        void $ sendx ki sock [ "on",L8.pack (show sz),bound ]
        void $ recvx ki sock sz
        return (EncSock sock ki sz bound,a)
      else do
        _ <- sendx ki sock ["no"]
        close sock
        wait4 asock


data EncSock = EncSock Socket L8.ByteString Int64 L8.ByteString
instance SockEnc EncSock where
  sk        (EncSock x _ _ _) = x
  nckey     (EncSock _ x _ _) = x
  chunksize (EncSock _ _ x _) = x
  bnd       (EncSock _ _ _ x) = x
instance Show EncSock where show _ = "undefined" -- TODO