DAmnProxy.hs

module Main where

import System.IO
import Network
import Data.ByteString.Lazy.Char8 (ByteString)
import Data.Char (digitToInt, ord, isHexDigit)
import Data.Unique
import Control.Monad (forever, liftM)
import Control.Concurrent (forkIO)
import Control.Exception (finally)
import Numeric (showHex)
import Data.Time.LocalTime (getZonedTime)
import Data.Time.Format (formatTime, defaultTimeLocale)
import qualified Data.ByteString.Lazy.Char8 as B
import qualified Network.WebSockets as WS

-- | Accept websocket connections.
acceptFromWS :: WS.ServerApp
acceptFromWS pendingConn = do
  putStrLn "Received connection!"
  conn <- WS.acceptRequest pendingConn
  dAmn <- makeDAmnConnection
  proxyToDAmn conn dAmn

-- | Proxy between dAmn and the websocket connection.
proxyToDAmn :: WS.Connection -> Handle -> IO ()
proxyToDAmn conn dAmn = do
  uuid <- newUnique
  let uuidHash = B.pack . show $ hashUnique uuid
      log msg = do
        t <- getZonedTime
        B.putStrLn $ "[" <+> B.pack (formatTime defaultTimeLocale "%r" t) <+> "]" <+> "(" <+> uuidHash <+> ") " <+> msg
  log "Starting proxying instance..."
  -- Writing to the websocket
  _ <- forkIO $ finally
    (forever $ readPackets dAmn >>= mapM_ (WS.sendTextData conn . urlEncode))
    (log ">> Done reading from dAmn. Disconnecting" >> hClose dAmn)
  
  -- Reading from the websocket
  forever (do
    msg <- WS.receiveData conn :: IO ByteString
    B.hPutStr dAmn $ urlDecode msg <+> "\0") `finally` (do
    log "<< Done reading from WS."
    WS.sendClose conn ("disconnect\ne=proxy error\n\0" :: ByteString)
    log "<< Closing dAmn..."
    hClose dAmn)
  
  log "End of proxy function"
  where inPacket pkt = B.putStrLn ("<in>" <+> pkt <+> "</in>") >> return pkt
        outPacket pkt = B.putStrLn ("<out>" <+> pkt <+> "</out>") >> return pkt

-- | Connect to dAmn.
makeDAmnConnection :: IO Handle
makeDAmnConnection = do
  putStrLn "Connecting to dAmn..."
  fd <- connectTo "chat.deviantart.com" $ PortNumber 3900
  putStrLn "Created dAmn socket."
  hSetBuffering fd NoBuffering
  hSetBinaryMode fd True -- if you don't do this then hGetContents will lag
  return fd

-- | Read a series of dAmn packets.
readPackets :: Handle -> IO [ByteString]
readPackets fd = liftM (B.split '\0') $ B.hGetContents fd

-- | URL decode the text.
urlDecode :: ByteString -> ByteString
urlDecode str =
  case B.break ('%'==) str of
    (a,"") -> a
    (a,b) ->
      let (code,rest) = B.splitAt 2 $ B.drop 1 b
          fstDigit = B.head code
          sndDigit = B.head $ B.tail code
          char = toEnum (16 * digitToInt fstDigit + digitToInt sndDigit)
      in if and [B.length code == 2, isHexDigit fstDigit, isHexDigit sndDigit]
            then a <+> (char `B.cons` urlDecode rest)
            else let (x,xs) = B.splitAt 1 b 
                 in a <+> x <+> urlDecode xs

-- | URL decode the text.
urlEncode :: ByteString -> ByteString
urlEncode str = -- gonna not bother with the space; must remember to escape %s so you don't mistakenly encode some garbage
  case B.break (\c -> c /= ' ' && c < '!' || c == '%' || c == '+' || c > '~') str of
    (a,"") -> a
    (a,b) ->
      let (char,rest) = B.splitAt 1 b
          code = B.pack $ '%' : (if (ord $ B.head char) <= 16 then '0' : hex else hex)
          hex = showHex (ord $ B.head char) ""
      in a <+> code <+> urlEncode rest

-- This should save me enough typing, no?
(<+>) = B.append

main = do
  putStrLn "Starting server..."
  WS.runServer "0.0.0.0" 9001 acceptFromWS