<commit_before>{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}

module Shadowsocks.Util
  ( Config (..)
  , cryptConduit
  , parseConfigOptions
  , unpackRequest
  , packRequest
  , packSockAddr
  , showSockAddr
  , SSException(..)
  ) where

import           Conduit (Conduit, awaitForever, yield, liftIO)
import           Control.Exception (Exception)
import           Control.Monad (liftM)
import           Data.Aeson (decode', FromJSON)
import           Data.Binary (decode)
import           Data.Binary.Get (runGet, getWord16be, getWord32le)
import           Data.Binary.Put (runPut, putWord16be, putWord32le)
import           Data.ByteString (ByteString)
import qualified Data.ByteString as S
import qualified Data.ByteString.Lazy as L
import qualified Data.ByteString.Char8 as C
import           Data.Char (chr, ord)
import           Data.IP ( fromHostAddress, fromHostAddress6
                         , toHostAddress, toHostAddress6)
import           Data.Maybe (fromMaybe)
import           Data.Typeable (Typeable)
import           GHC.Generics (Generic)
import           Network.Socket (HostAddress, HostAddress6, SockAddr(..))
import           Options.Applicative
import           System.Directory (doesFileExist)
import           System.Exit (exitFailure)


data Config = Config
    { server        :: String
    , server_port   :: Int
    , local_port    :: Int
    , password      :: String
    , timeout       :: Int
    , method        :: String
    } deriving (Show, Generic)

instance FromJSON Config

data Options = Options
    { _server        :: Maybe String
    , _server_port   :: Maybe Int
    , _local_port    :: Maybe Int
    , _password      :: Maybe String
    , _method        :: Maybe String
    , _config        :: Maybe String
    } deriving (Show, Generic)

type AddrType = Int

data SSException = UnknownAddrType AddrType 
                 | NoRequestBody
    deriving (Show, Typeable)

instance Exception SSException


nullConfig :: Config
nullConfig = Config "" 0 0 "" 0 ""

readConfig :: FilePath -> IO (Maybe Config)
readConfig fp = liftM decode' $ L.readFile fp

configOptions :: Parser Options
configOptions = Options
    <$> optional (strOption (long "server" <> short 's' <> metavar "ADDR"
              <> help "server address"))
    <*> optional (option auto (long "server-port" <> short 'p' <> metavar "PORT"
                <> help "server port"))
    <*> optional (option auto (long "local-port" <> short 'l' <> metavar "PORT"
                <> help "local port"))
    <*> optional (strOption (long "password" <> short 'k' <> metavar "PASSWORD"
              <> help "password"))
    <*> optional (strOption (long "method" <> short 'm' <> metavar "METHOD"
              <> help "encryption method, for example, aes-256-cfb"))
    <*> optional (strOption (long "config" <> short 'c' <> metavar "CONFIG"
               <> help "path to config file"))

parseConfigOptions :: IO Config
parseConfigOptions = do
    o <- execParser $ info (helper <*> configOptions)
                      (fullDesc <> header "shadowsocks - a fast tunnel proxy")
    let configFile = fromMaybe "config.json" (_config o)
    doesFileExist configFile >>= \case
        True -> putStrLn "INFO: config file loaded"
        False -> putStrLn "ERROR: config file not found" >> exitFailure
    mconfig <- readConfig configFile
    let c = fromMaybe nullConfig mconfig
    return $ c { server = fromMaybe (server c) (_server o)
               , server_port = fromMaybe (server_port c) (_server_port o)
               , local_port = fromMaybe (local_port c) (_local_port o)
               , password = fromMaybe (password c) (_password o)
               , method = fromMaybe (method c) (_method o)
               }

cryptConduit :: (ByteString -> IO ByteString)
             -> Conduit ByteString IO ByteString
cryptConduit crypt = awaitForever $ \input -> do
    output <- liftIO $ crypt input
    yield output

unpackRequest :: ByteString -> Either AddrType (AddrType, ByteString, Int, ByteString)
unpackRequest request = case addrType of
    1 ->        -- IPv4
        let (ip, rest) = S.splitAt 4 request'
            addr = C.pack $ show $ fromHostAddress $ runGet getWord32le
                                                   $ L.fromStrict ip
        in  Right (addrType, addr, unpackPort rest, S.drop 2 rest)
    3 ->        -- domain name
        let addrLen = ord $ C.head request'
            (domain, rest) = S.splitAt (addrLen + 1) request'
        in  Right (addrType, S.tail domain, unpackPort rest, S.drop 2 rest)
    4 ->        -- IPv6
        let (ip, rest) = S.splitAt 16 request'
            addr = C.pack $ show $ fromHostAddress6 $ decode
                                                    $ L.fromStrict ip
        in  Right (addrType, addr, unpackPort rest, S.drop 2 rest)<commit_msg>catch IOException instead<commit_after>{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Shadowsocks.Util
  ( Config (..)
  , cryptConduit
  , parseConfigOptions
  , unpackRequest
  , packRequest
  , packSockAddr
  , showSockAddr
  , SSException(..)
  ) where

import           Conduit (Conduit, awaitForever, yield, liftIO)
import           Control.Exception (Exception, IOException, catch)
import           Control.Monad (liftM)
import           Data.Aeson (decode', FromJSON)
import           Data.Binary (decode)
import           Data.Binary.Get (runGet, getWord16be, getWord32le)
import           Data.Binary.Put (runPut, putWord16be, putWord32le)
import           Data.ByteString (ByteString)
import qualified Data.ByteString as S
import qualified Data.ByteString.Lazy as L
import qualified Data.ByteString.Char8 as C
import           Data.Char (chr, ord)
import           Data.IP ( fromHostAddress, fromHostAddress6
                         , toHostAddress, toHostAddress6)
import           Data.Maybe (fromMaybe)
import           Data.Typeable (Typeable)
import           GHC.Generics (Generic)
import           Network.Socket (HostAddress, HostAddress6, SockAddr(..))
import           Options.Applicative
import           System.Exit (exitFailure)
import           System.IO (hPutStrLn, stderr)


data Config = Config
    { server        :: String
    , server_port   :: Int
    , local_port    :: Int
    , password      :: String
    , timeout       :: Int
    , method        :: String
    } deriving (Show, Generic)

instance FromJSON Config

data Options = Options
    { _server        :: Maybe String
    , _server_port   :: Maybe Int
    , _local_port    :: Maybe Int
    , _password      :: Maybe String
    , _method        :: Maybe String
    , _config        :: Maybe String
    } deriving (Show, Generic)

type AddrType = Int

data SSException = UnknownAddrType AddrType 
                 | NoRequestBody
    deriving (Show, Typeable)

instance Exception SSException


nullConfig :: Config
nullConfig = Config "" 0 0 "" 0 ""

readConfig :: FilePath -> IO (Maybe Config)
readConfig fp = liftM decode' $ L.readFile fp

configOptions :: Parser Options
configOptions = Options
    <$> optional (strOption (long "server" <> short 's' <> metavar "ADDR"
              <> help "server address"))
    <*> optional (option auto (long "server-port" <> short 'p' <> metavar "PORT"
                <> help "server port"))
    <*> optional (option auto (long "local-port" <> short 'l' <> metavar "PORT"
                <> help "local port"))
    <*> optional (strOption (long "password" <> short 'k' <> metavar "PASSWORD"
              <> help "password"))
    <*> optional (strOption (long "method" <> short 'm' <> metavar "METHOD"
              <> help "encryption method, for example, aes-256-cfb"))
    <*> optional (strOption (long "config" <> short 'c' <> metavar "CONFIG"
               <> help "path to config file"))

parseConfigOptions :: IO Config
parseConfigOptions = do
    o <- execParser $ info (helper <*> configOptions)
                      (fullDesc <> header "shadowsocks - a fast tunnel proxy")
    let configFile = fromMaybe "config.json" (_config o)
    mconfig <- readConfig configFile `catch` \(e :: IOException) ->
        hPutStrLn stderr ("ERROR: Failed to load " <> show e) >> exitFailure
    let c = fromMaybe nullConfig mconfig
    return $ c { server = fromMaybe (server c) (_server o)
               , server_port = fromMaybe (server_port c) (_server_port o)
               , local_port = fromMaybe (local_port c) (_local_port o)
               , password = fromMaybe (password c) (_password o)
               , method = fromMaybe (method c) (_method o)
               }

cryptConduit :: (ByteString -> IO ByteString)
             -> Conduit ByteString IO ByteString
cryptConduit crypt = awaitForever $ \input -> do
    output <- liftIO $ crypt input
    yield output

unpackRequest :: ByteString -> Either AddrType (AddrType, ByteString, Int, ByteString)
unpackRequest request = case addrType of
    1 ->        -- IPv4
        let (ip, rest) = S.splitAt 4 request'
            addr = C.pack $ show $ fromHostAddress $ runGet getWord32le
                                                   $ L.fromStrict ip
        in  Right (addrType, addr, unpackPort rest, S.drop 2 rest)
    3 ->        -- domain name
        let addrLen = ord $ C.head request'
            (domain, rest) = S.splitAt (addrLen + 1) request'
        in  Right (addrType, S.tail domain, unpackPort rest, S.drop 2 rest)
    4 ->        -- IPv6
        let (ip, rest) = S.splitAt 16 request'
            addr = C.pack $ show $ fromHostAddress6 $ decode
                                                    $ L.fromStrict ip
        in  Right (addrType, addr, unpackPort rest, S.drop 2 rest)
{'fraction_non_alphanumeric': 0.10170356013465266, 'fraction_numerical': 0.006732632867489544, 'mean_word_length': 2.234576047509073, 'pattern_counts': {'":': 0, '<': 62, '<?xml version=': 0, '>': 81, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3184354', 'n_tokens_mistral': 2858, 'n_tokens_neox': 2734, 'n_words': 875}