fradrive/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
  , sourceFiles, acceptFile
  ) where

import Import

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

-- import qualified Data.ByteString.Lazy as Lazy (ByteString)
import qualified Data.ByteString.Lazy as Lazy.ByteString

import qualified Data.ByteString as ByteString

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


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.
              ( MonadThrow b
              , MonadThrow m
              , MonadBase b m
              , PrimMonad b
              )
           => ConduitT ByteString File m ZipInfo
consumeZip = transPipe liftBase unZipStream `fuseUpstream` consumeZip'
  where
    consumeZip' :: ConduitT (Either ZipEntry ByteString) File m ()
    consumeZip' = do
      input <- await
      case input of
        Nothing -> return ()
        Just (Right _) -> throwM $ userError "Data chunk in unexpected place when parsing ZIP"
        Just (Left ZipEntry{..}) -> do
          contentChunks <- toConsumer accContents
          zipEntryName' <- decodeZipEntryName zipEntryName
          let
            fileTitle = dropWhile isPathSeparator . dropTrailingPathSeparator . normalise $ makeValid zipEntryName'
            fileModified = localTimeToUTC utc zipEntryTime
            fileContent
              | hasTrailingPathSeparator zipEntryName' = Nothing
              | otherwise = Just $ mconcat contentChunks
          yield File{..}
          consumeZip'
    accContents :: ConduitT (Either a b') Void m [b']
    accContents = do
      input <- await
      case input of
        Just (Right x) -> (x :) <$> accContents
        Just (Left x)  -> []    <$ leftover (Left x)
        _ -> return []

produceZip :: forall b m.
            ( MonadThrow b
            , MonadThrow m
            , MonadBase b m
            , PrimMonad b
            )
           => ZipInfo
           -> ConduitT File ByteString m ()
produceZip info = C.map toZipData .| transPipe liftBase (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 -> (ZipEntry, ZipData b)
    toZipData f@File{..} =
      let zData = maybe mempty (ZipDataByteString . Lazy.ByteString.fromStrict) fileContent
          zEntry = (toZipEntry f){ zipEntrySize = fromIntegral . ByteString.length <$> fileContent }
       in (zEntry, zData)

    toZipEntry :: File -> 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 File m ()
modifyFileTitle f = mapC $ \x@File{..} -> x{ fileTitle = f fileTitle }

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

acceptFile :: MonadResource m => FileInfo -> m File
acceptFile fInfo = do
  let fileTitle = dropWhile isPathSeparator . dropTrailingPathSeparator . normalise . unpack $ fileName fInfo
  fileModified <- liftIO getCurrentTime
  fileContent <- fmap Just . runConduit $ fileSource fInfo .| foldC
  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