fradrive-old/src/Handler/Utils/Zip.hs

{-# OPTIONS_GHC -fno-warn-missing-fields #-} -- This concerns zipEntrySize in produceZip
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Handler.Utils.Zip
  ( typeZip, extensionZip
  , ZipError(..)
  , ZipInfo(..)
  , produceZip
  , consumeZip
  , modifyFileTitle
  , receiveFiles, acceptFile
  ) where

import Import

import Codec.Archive.Zip.Conduit.Types
import Codec.Archive.Zip.Conduit.UnZip
import Codec.Archive.Zip.Conduit.Zip

import System.FilePath
import Data.Time.LocalTime (localTimeToUTC, utcToLocalTime)

import Data.List (dropWhileEnd)

import qualified Data.Map as Map

import qualified Data.Conduit.List as C


import Data.Encoding ( decodeStrictByteStringExplicit
                     , encodeStrictByteStringExplicit
                     )
import Data.Encoding.CP437
import qualified Data.Char as Char

import Control.Monad.Trans.Cont


typeZip :: ContentType
typeZip = "application/zip"

extensionZip :: Extension
extensionZip = fromMaybe "zip" $ listToMaybe [ ext | (ext, mime) <- Map.toList mimeMap, mime == typeZip ]



instance Default ZipInfo where
  def = ZipInfo
    { zipComment = mempty
    }


consumeZip :: forall b m m'.
              ( MonadThrow b
              , MonadThrow m
              , MonadBase b m
              , PrimMonad b
              , MonadUnliftIO m
              , MonadResource m
              , MonadIO m'
              )
           => ConduitT () ByteString m () -> ConduitT () (File m') m ZipInfo
consumeZip inpBS = do
  inps <- liftIO newBroadcastTMChanIO
  let feedSingle inp = atomically $ do
        guardM $ isEmptyTMChan inps
        writeTMChan inps inp
  zipAsync <- lift . allocateLinkedAsync . runConduit $ do
    zipInfo <- (inpBS .| transPipe liftBase unZipStream) `fuseUpstream` C.mapM_ feedSingle
    atomically $ closeTMChan inps
    return zipInfo

  evalContT . callCC $ \finish -> forever $ do
    (fileChan, fileDef) <- atomically $ do
      fileChan <- dupTMChan inps
      fileDef <- readTMChan fileChan
      return (fileChan, fileDef)
    case fileDef of
      Nothing -> finish =<< waitAsync zipAsync
      Just (Right _) -> return ()
      Just (Left ZipEntry{..}) -> do
        zipEntryName' <- decodeZipEntryName zipEntryName
        let
          fileTitle = dropWhile isPathSeparator . dropTrailingPathSeparator . normalise $ makeValid zipEntryName'
          fileModified = localTimeToUTC utc zipEntryTime
          isDirectory = hasTrailingPathSeparator zipEntryName'
          fileContent
            | isDirectory = Nothing
            | otherwise   = Just . evalContT . callCC $ \finishContent -> forever $ do
                nextVal <- atomically $ (preview _Right =<<) <$> readTMChan fileChan
                maybe (finishContent ()) (lift . yield) nextVal
        lift $ yield File{..}

produceZip :: forall m.
            ( MonadThrow m
            , PrimMonad m
            )
           => ZipInfo
           -> ConduitT (File m) ByteString m ()
produceZip info = C.map toZipData .| void (zipStream zipOptions)
  where
    zipOptions = ZipOptions
      { zipOpt64 = True
      , zipOptCompressLevel = -1 -- This is passed through all the way to the C zlib, where it means "default level"
      , zipOptInfo = info
      }

    toZipData :: File m -> (ZipEntry, ZipData m)
    toZipData f@File{..}
       = (toZipEntry f, ) $ maybe mempty ZipDataSource fileContent

    toZipEntry :: File m -> ZipEntry
    toZipEntry File{..} = ZipEntry{..}
      where
        isDir = isNothing fileContent

        zipEntryName = encodeZipEntryName . bool (dropWhileEnd isPathSeparator) addTrailingPathSeparator isDir . normalise $ makeValid fileTitle
        zipEntryTime = utcToLocalTime utc fileModified
        zipEntrySize = Nothing
        zipEntryExternalAttributes = Nothing

modifyFileTitle :: Monad m => (FilePath -> FilePath) -> ConduitT (File m') (File m') m ()
modifyFileTitle f = mapC $ \x@File{..} -> x{ fileTitle = f fileTitle }

-- Takes FileInfo and if it is a ZIP-Archive, extract files, otherwiese yield fileinfo
receiveFiles :: (MonadLogger m, MonadResource m, MonadThrow m, MonadBase IO m, MonadUnliftIO m, MonadResource m') => FileInfo -> ConduitT () (File m') m ()
receiveFiles fInfo
    | ((==) `on` simpleContentType) mimeType typeZip = do
        $logInfoS "sourceFiles" "Unpacking ZIP"
        void . consumeZip $ fileSource fInfo
    | otherwise = do
        $logDebugS "sourceFiles" [st|Not unpacking file of type #{decodeUtf8 mimeType}|]
        yieldM $ acceptFile fInfo
  where
    mimeType = mimeLookup $ fileName fInfo

acceptFile :: (MonadResource m, MonadResource m') => FileInfo -> m (File m')
acceptFile fInfo = do
  let fileTitle = dropWhile isPathSeparator . dropTrailingPathSeparator . normalise . unpack $ fileName fInfo
      fileContent = Just $ fileSource fInfo
  fileModified <- liftIO getCurrentTime
  return File{..}


decodeZipEntryName :: MonadThrow m => Either Text ByteString -> m FilePath
-- ^ Extract the filename from a 'ZipEntry' doing decoding along the way.
--
-- Throws 'Data.Encoding.Exception.DecodingException's.
decodeZipEntryName = \case
  Left  t     -> return $ unpack t
  Right cp437 -> either throwM return $ decodeStrictByteStringExplicit CP437 cp437

encodeZipEntryName :: FilePath -> Either Text ByteString
-- ^ Encode a filename for use in a 'ZipEntry', encodes as
--   'Data.Encoding.UTF8.UTF8' iff the given path contains non-ascii characters.
--
-- Does not do any normalisation (in particular this function does not ensure
-- that the 'FilePath' does not start with a slash).
encodeZipEntryName path = fromMaybe (Left $ pack path) $ do
  guard $ all Char.isAscii path

  either (const mzero) (return . Right) $ encodeStrictByteStringExplicit CP437 path