diff --git a/src/Config.hs b/src/Config.hs new file mode 100644 index 0000000..5904c76 --- /dev/null +++ b/src/Config.hs @@ -0,0 +1,10 @@ +module Config where + +import System.Directory + +data ViddlConfig = ViddlConfig + { webPort :: Int + , dlDir :: FilePath } + +defConfig :: IO ViddlConfig +defConfig = getTemporaryDirectory >>= pure . ViddlConfig 3000 . (<> "/viddl") diff --git a/src/Helpers.hs b/src/Helpers.hs index 7d666f6..19b4589 100644 --- a/src/Helpers.hs +++ b/src/Helpers.hs @@ -3,8 +3,26 @@ module Helpers where import YTDL +import Config + +import Text.Read (readMaybe) import qualified Data.Text.Lazy as TL -import Network.URI (parseURI) +import Network.URI +import System.Exit + +parseArgs :: ViddlConfig -> [String] -> IO ViddlConfig +parseArgs _ ("-h":_) = exitHelp +parseArgs _ ("--help":_) = exitHelp +parseArgs cfg@(ViddlConfig _ d) ("-p":p':xs) = case (readMaybe p' :: Maybe Int) of + (Just p) -> parseArgs (ViddlConfig p d) xs + Nothing -> parseArgs cfg (p':xs) +parseArgs (ViddlConfig p _) ("-d":d:xs) = parseArgs (ViddlConfig p d) xs +parseArgs cfg (_:xs) = parseArgs cfg xs +parseArgs cfg [] = pure cfg + + +exitHelp :: IO a +exitHelp = putStrLn "viddl [-p port] [-d dir]" >> exitSuccess getRes :: TL.Text -> Maybe Resolution getRes ("144p") = Just P144 @@ -22,7 +40,13 @@ isRes res = case getRes res of (Just _) -> True _ -> False +-- ssrf paranoia +isOkPath :: String -> Bool +isOkPath p = not $ isIPv4address p || isIPv6address p || p == "localhost" + isURL :: TL.Text -> Bool isURL uri = case parseURI (TL.unpack uri) of - (Just _) -> True - _ -> False + (Just u) -> case uriAuthority u of + (Just (URIAuth _ p _)) -> isOkPath p + _ -> False + _ -> False diff --git a/src/Main.hs b/src/Main.hs index 3d5a504..474eead 100644 --- a/src/Main.hs +++ b/src/Main.hs @@ -6,37 +6,46 @@ import Templates.Index import Templates.Error import YTDL import Helpers +import Config -import Control.Monad.IO.Class +import Control.Monad.Trans.Reader +import Control.Monad.Trans.Class import qualified Data.Text.Lazy as TL -import Web.Scotty +import System.Environment +import Web.Scotty.Trans -safeDownloadAction :: ActionM () +type Scotty = ScottyT TL.Text (ReaderT ViddlConfig IO) +type Action = ActionT TL.Text (ReaderT ViddlConfig IO) + +safeDownloadAction :: Action () safeDownloadAction = downloadAction `rescue` (html . errorPage) -downloadAction :: ActionM () +downloadAction :: Action () downloadAction = do url <- param "url" res <- param "resolution" if (isURL url) && (isRes res) then do let (Just res') = getRes res -- safe cause we checked with isRes - ytdlRes <- liftIO $ ytdl (TL.unpack url) res' + ytdlRes <- lift $ ytdl (TL.unpack url) res' case ytdlRes of (Right filePath) -> do setHeader "content-type" "video/mp4" file filePath (Left err) -> html $ errorPage (TL.pack err) - -- liftIO $ ytdlClean ident else html $ errorPage "Invalid input!" -- todo ReaderT -app :: ScottyM () +app :: Scotty () app = do get "/" $ html indexPage get "/video.mp4" safeDownloadAction get "/audio.mp3" safeDownloadAction main :: IO () -main = scotty 3000 app +main = do + args <- getArgs + defCfg <- defConfig + cfg <- parseArgs defCfg args + scottyT (webPort cfg) (\(ReaderT ma) -> ma cfg) app diff --git a/src/Templates/Error.hs b/src/Templates/Error.hs index 47ae1fd..7356fa4 100644 --- a/src/Templates/Error.hs +++ b/src/Templates/Error.hs @@ -5,6 +5,7 @@ module Templates.Error (errorPage) where import qualified Data.Text.Lazy as TL import Text.RawString.QQ +import Text.HTML.SanitizeXSS errorPage :: TL.Text -> TL.Text errorPage msg = [r| @@ -17,7 +18,7 @@ errorPage msg = [r|

viddl

-

|] <> msg <> [r|

+

|] <> (TL.fromStrict (sanitize (TL.toStrict msg))) <> [r|

viddl is free open source software and is powered by youtube-dl.

diff --git a/src/YTDL.hs b/src/YTDL.hs index 8f5375a..85d84c2 100644 --- a/src/YTDL.hs +++ b/src/YTDL.hs @@ -1,5 +1,8 @@ module YTDL (Resolution(..), ytdl) where +import Config + +import Control.Monad.Trans.Reader import Control.Concurrent (forkIO, threadDelay) import Data.Digest.Pure.MD5 import qualified Data.ByteString.Lazy.UTF8 as BCU @@ -21,7 +24,7 @@ data Resolution wrapResString :: String -> [String] wrapResString str = ["-f", concat ["bestvideo[ext=mp4,height<=", str, - "]+bestaudio[ext=m4a]/mp4[height<=", str, "]"]] + "]+bestaudio[ext=m4a]/mp4[height<=", str, "]"], "--merge-output-format", "mp4"] resToArgs :: Resolution -> [String] resToArgs (P144) = wrapResString "144" @@ -30,21 +33,22 @@ resToArgs (P360) = wrapResString "360" resToArgs (P480) = wrapResString "480" resToArgs (P720) = wrapResString "720" resToArgs (P1080) = wrapResString "1080" -resToArgs (PMAX) = ["-f", "bestvideo[ext=mp4]+bestaudio[ext=m4a]/mp4"] +resToArgs (PMAX) = ["-f", "bestvideo[ext=mp4]+bestaudio[ext=m4a]/mp4", "--merge-output-format", "mp4"] resToArgs (Audio) = ["-x", "--audio-format", "mp3"] -ytdl :: String -> Resolution -> IO (Either String FilePath) -ytdl url res = do +extraYtdlArgs :: [String] +extraYtdlArgs = ["--no-playlist"] + +ytdl :: String -> Resolution -> ReaderT ViddlConfig IO (Either String FilePath) +ytdl url res = ReaderT $ \cfg -> do + let ext = case res of { Audio -> ".mp3"; _ -> ".mp4" } let ident = show . md5 . BCU.fromString $ url <> show res - -- todo: config for path - tmpdir <- getTemporaryDirectory - let dir = concat [tmpdir, "/viddl/", ident] + let dir = concat [dlDir cfg, "/", ident] let fileName = concat [dir, "/", ident, ext] - -- todo: implement file locking for deleting and uploading file processed <- doesFileExist fileName if processed then do @@ -55,13 +59,16 @@ ytdl url res = do createDirectoryIfMissing True dir - ytdlProc <- createProcess (proc "youtube-dl" (resToArgs res <> ["-o", fileName, url])) + print (resToArgs res <> ["-o", fileName, url]) + + ytdlProc <- createProcess (proc "youtube-dl" (resToArgs res <> ["-o", fileName, url] <> extraYtdlArgs)) { std_out = CreatePipe , std_err = CreatePipe } case ytdlProc of - (_, _, Just herr, ph) -> do + (_, Just hout, Just herr, ph) -> do err <- hGetContents herr + out <- hGetContents hout exitCode <- waitForProcess ph case exitCode of ExitSuccess -> do @@ -69,7 +76,6 @@ ytdl url res = do if exists then do -- wait 5 minutes and then delete the directory - -- todo: implement file locking for deleting and uploading file _ <- forkIO $ threadDelay 300000000 >> removeDirectoryRecursive dir pure (Right fileName) else do diff --git a/todo b/todo index 131fce8..9a2b614 100644 --- a/todo +++ b/todo @@ -1,2 +1,3 @@ -configuration for scotty -cache results somehow +better config options (file?) +safer file deletion +queue system? diff --git a/viddl.cabal b/viddl.cabal index a994d7b..3528ae3 100644 --- a/viddl.cabal +++ b/viddl.cabal @@ -17,8 +17,9 @@ library , Templates.Error , YTDL , Helpers + , Config ghc-options: -Wall -O2 -threaded - build-depends: base ^>=4.14.1.0 + build-depends: base , scotty , text , raw-strings-qq @@ -28,6 +29,8 @@ library , bytestring , pureMD5 , utf8-string + , transformers + , xss-sanitize hs-source-dirs: src default-language: Haskell2010 @@ -37,8 +40,9 @@ executable viddl , Templates.Error , Helpers , YTDL + , Config ghc-options: -Wall -O2 -threaded - build-depends: base ^>=4.14.1.0 + build-depends: base , scotty , text , raw-strings-qq @@ -48,5 +52,7 @@ executable viddl , bytestring , pureMD5 , utf8-string + , transformers + , xss-sanitize hs-source-dirs: src default-language: Haskell2010