wsedit-0.1.0.0: wsedit.hs

--
--  Copyright (C) 2012 Linus Lüssing
--
--  This program is free software: you can redistribute it and/or modify
--  it under the terms of the GNU General Public License as published by
--  the Free Software Foundation, either version 3 of the License, or
--  (at your option) any later version.
--
--  This program is distributed in the hope that it will be useful,
--  but WITHOUT ANY WARRANTY; without even the implied warranty of
--  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
--  GNU General Public License for more details.
--
--  You should have received a copy of the GNU General Public License
--  along with this program.  If not, see <http://www.gnu.org/licenses/>.
--

-- TODO: Somehow, partially merge with / export to Data.Torrent code

import Control.Applicative
import Control.Monad

import System.Environment
import System.Exit
import System.IO
import System.Directory

import qualified Data.Map as M
import qualified Data.List as L
import Data.BEncode

import Safe (readMay)

import qualified Data.ByteString.Lazy as B
import qualified Data.ByteString.Lazy.UTF8 as BUTF8

---

-- TODO: Make this safe, replace 'fail' with Either or Maybe
parseTorrent' :: IO (Maybe BEncode) -> IO BEncode
parseTorrent' ioMaybeBencode = do
	maybeBencode <- ioMaybeBencode
	case maybeBencode of
		Nothing -> fail "Error, could not parse file"
		(Just bencode) -> return bencode

parseTorrent :: String -> IO BEncode
parseTorrent file = do
	let contents = B.readFile file
	parseTorrent' $ fmap bRead contents

---

filterStringsFromBList :: [BEncode] -> [BEncode]
filterStringsFromBList =
	filter (\x -> case x of
			(BString _) -> True
			otherwise -> False)

webseedsFromBList :: [BEncode] -> [String]
webseedsFromBList = map toString . filterStringsFromBList
	where toString (BString x) = BUTF8.toString x

webseedsFromKey :: BEncode -> [String]
webseedsFromKey urls =
	case urls of
		BString url -> map BUTF8.toString [url]
		BList list -> webseedsFromBList list
		otherwise -> []

webseedsFromBDict :: M.Map String BEncode -> [String]
webseedsFromBDict dict =
	let maybeUrls = M.lookup "url-list" dict in
	case maybeUrls of
		Nothing -> []
		(Just urls) -> webseedsFromKey urls

getWebseeds :: BEncode -> [String]
getWebseeds bencode =
	case bencode of
		BDict dict -> appendNums $ webseedsFromBDict dict
		otherwise -> []
	where appendNums = zipWith (\n line -> show n ++ " - " ++ line) [0..]

---

exitWithErrorMsg :: String -> IO ()
exitWithErrorMsg err = do
	putStrLn (err ++ "\n")
	printHelp ()
	exitWith (ExitFailure 1)

exitInvalidArgument :: IO ()
exitInvalidArgument = exitWithErrorMsg "Error, invalid argument!"

writeBEncode :: String -> BEncode -> IO ()
writeBEncode file bencode = do
	(tempName, tempHandle) <- openBinaryTempFile "." "temp"
	let stream = bPack bencode
	B.hPut tempHandle stream
	hClose tempHandle
	renameFile tempName file

---

listWebseeds :: [String] -> IO ()
listWebseeds [] = exitInvalidArgument
listWebseeds [file] = fmap getWebseeds (parseTorrent file) >>= mapM_ putStrLn

---

stringToBString :: String -> BEncode
stringToBString = BString . BUTF8.fromString

appendUrl :: String -> BEncode -> BEncode
appendUrl url (BList list) = BList (list ++ [stringToBString url])

-- Case: There's no url-list in the dictionary yet, we need to create
-- one as a simple bencoded list
addFirstWebseed :: String -> M.Map String BEncode -> BEncode
addFirstWebseed url dict = BDict $ M.insert "url-list" (stringToBString url) dict

-- Case: There's already a single bencoded string, we need to convert it to
-- a list and append our new url afterwards
addSecondWebseed :: String -> M.Map String BEncode -> BEncode
addSecondWebseed url dict = BDict $ M.adjust (appendUrl url . toBList) "url-list" dict
	where toBList x = BList [x]

-- Case: There's already a bencoded list which contains at least two
-- url strings
addAnotherWebseed :: String -> M.Map String BEncode -> BEncode
addAnotherWebseed url dict = BDict $ M.adjust (appendUrl url) "url-list" dict


addWebseed' :: String -> BEncode -> Either String BEncode
addWebseed' newurl (BDict dict) =
	let maybeUrls = M.lookup "url-list" dict in
	case maybeUrls of
		Nothing -> Right (addFirstWebseed newurl dict)
		(Just (BString url)) -> Right (addSecondWebseed newurl dict)
		(Just (BList urls)) -> Right (addAnotherWebseed newurl dict)
		otherwise -> Left "Error, incompatible (broken?) url-list found!"
addWebseed' _ _ = Left "Error, not a valid torrent file!"

addWebseed :: [String] -> IO ()
addWebseed [] = exitInvalidArgument
addWebseed [x] = exitInvalidArgument
addWebseed [url,file] = do
	newBencoded <- fmap (addWebseed' url) $ parseTorrent file
	case newBencoded of
		(Right y) -> writeBEncode file y
		(Left y) -> exitWithErrorMsg y
addWebseed _ = exitInvalidArgument

---

removeWebseed''' :: Integer -> BEncode -> BEncode
removeWebseed''' num (BList [item0,item1]) = if num == 0 then item1 else item0
removeWebseed''' num (BList list) = BList $ remove num list
	where remove n xs = let (ys,zs) = L.genericSplitAt n xs in ys ++ tail zs

removeWebseed'' :: Integer -> BEncode -> Either String BEncode
removeWebseed'' num bencoded@(BDict dict) =
	let maybeUrls = M.lookup "url-list" dict in
	case maybeUrls of
		Nothing -> Left "Error, the torrent has no webseeds!"
		(Just (BString url)) -> if num > 0
						then errTooLargeNum
						else Right (BDict $ M.delete "url-list" dict)
		(Just (BList urls)) -> if fromIntegral (length urls) < (num + 1)
						then errTooLargeNum
						else Right (BDict $ M.adjust (removeWebseed''' num) "url-list" dict)
	where	errTooLargeNum = Left ("Error, number " ++ show num ++ " is too large! " ++ urlStrList)
		urlStrList = "Choose:\n\n" ++ init (unlines $ getWebseeds bencoded)
removeWebseed'' _ _ = Left "Error, not a valid torrent file!"

removeWebseed' :: Integer -> String -> IO ()
removeWebseed' num file = do
	newBencoded <- fmap (removeWebseed'' num) $ parseTorrent file
	case newBencoded of
		(Right y) -> writeBEncode file y
		(Left y) -> exitWithErrorMsg y

removeWebseed :: [String] -> IO ()
removeWebseed [] = exitInvalidArgument
removeWebseed [x] = exitInvalidArgument
removeWebseed [num,file] =
	let maybeNum = readMay num :: Maybe Integer in
	case maybeNum of
		(Just i) -> removeWebseed' i file
		Nothing -> exitWithErrorMsg "Error, input is not a valid integer number!"
removeWebseed _ = exitInvalidArgument


removeAllWebseeds' :: BEncode -> BEncode
removeAllWebseeds' bencode = case bencode of
				BDict dict -> BDict $ M.delete "url-list" dict
				otherwise -> bencode

removeAllWebseeds :: [String] -> IO ()
removeAllWebsedes [] = exitInvalidArgument
removeAllWebseeds [file] = fmap removeAllWebseeds' (parseTorrent file) >>= writeBEncode file
removeAllWebseeds _ = exitInvalidArgument

---

-- TODO: Change to System.Console.GetOpt/CmdArgs/CmdLib/ParseArgs/...
printHelp :: a -> IO ()
printHelp _ = do
	progName <- getProgName
	putStrLn ("Usage: " ++ progName ++ " [options] torrentfile")
	putStrLn ("   or: " ++ progName ++ " --help|-h\n")
	putStrLn "Options:"
	putStrLn "\t--list-webseeds\t\tDisplay webseeds"
	putStrLn "\t--add-webseed <url>\tAppend new webseed to torrent file"
	putStrLn "\t--remove-webseed <n>\tRemove the n-th webseed, starting with 0"
	putStrLn "\t--remove-all-webseeds\tRemove all webseeds from torrent file"

dispatch :: [(String, [String] -> IO ())]
dispatch =
	[ ("--help", printHelp)
	, ("-h", printHelp)
	, ("--list-webseeds", listWebseeds)
	, ("--add-webseed", addWebseed)
	, ("--remove-webseed", removeWebseed)
	, ("--remove-all-webseeds", removeAllWebseeds)
	]


execCommand :: String -> [String] -> IO ()
execCommand cmd args = do
	let maybeAction = L.lookup cmd dispatch
	case maybeAction of
		Nothing -> do	putStrLn "Error, invalid command\n"
				printHelp ()
				exitWith (ExitFailure 1)
		(Just action) -> action args

main = do
	args <- getArgs
	case args of
		[] -> printHelp ()
		command:cargs -> execCommand command cargs