-- hot.hs: hOpenPGP Tool
-- Copyright © 2012-2021  Clint Adams
--
-- vim: softtabstop=4:shiftwidth=4:expandtab
--
-- This program is free software: you can redistribute it and/or modify
-- it under the terms of the GNU Affero 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 Affero General Public License for more details.
--
-- You should have received a copy of the GNU Affero General Public License
-- along with this program.  If not, see <http://www.gnu.org/licenses/>.

{-# LANGUAGE RecordWildCards #-}

import qualified Codec.Encryption.OpenPGP.ASCIIArmor as AA
import Codec.Encryption.OpenPGP.ASCIIArmor.Types (Armor(..), ArmorType(..))
import Codec.Encryption.OpenPGP.Serialize ()
import Codec.Encryption.OpenPGP.Types
import Control.Applicative (optional)
import Control.Error.Util (note)
import Control.Monad.IO.Class (MonadIO, liftIO)
import qualified Data.Aeson as A
import Data.Binary (get, put)
import Data.Binary.Get (Get)
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import Data.Conduit (ConduitM, (.|), runConduitRes)
import qualified Data.Conduit.Binary as CB
import qualified Data.Conduit.List as CL
import Data.Conduit.OpenPGP.Filter
  ( FilterPredicates(RPFilterPredicate)
  , conduitPktFilter
  )
import Data.Conduit.Serialization.Binary (conduitGet, conduitPut)
import Data.Void (Void)
import qualified Data.Yaml as Y
import HOpenPGP.Tools.Armor (doDeArmor)
import HOpenPGP.Tools.Common (banner, prependAuto, versioner, warranty)
import HOpenPGP.Tools.Parser (parsePExp)

import System.IO
  ( BufferMode(..)
  , Handle
  , hFlush
  , hSetBuffering
  , stderr
  , stdin
  , stdout
  )

import Data.Text.Prettyprint.Convert.AnsiWlPprint (toAnsiWlPprint)
import Data.Text.Prettyprint.Doc
  ( Pretty
  , (<+>)
  , group
  , hardline
  , list
  , pretty
  , softline
  )
import Data.Text.Prettyprint.Doc.Render.Text (hPutDoc)
import Options.Applicative.Builder
  ( argument
  , command
  , eitherReader
  , footerDoc
  , headerDoc
  , help
  , helpDoc
  , info
  , long
  , metavar
  , option
  , prefs
  , progDesc
  , showDefaultWith
  , showHelpOnError
  , str
  , strOption
  , value
  )
import Options.Applicative.Extra (customExecParser, helper, hsubparser)
import Options.Applicative.Types (Parser)

data Command
  = DumpC DumpOptions
  | DeArmorC
  | ArmorC ArmoringOptions
  | FilterC FilteringOptions

data DumpOptions =
  DumpOptions
    { outputformat :: DumpOutputFormat
    }

data FilteringOptions =
  FilteringOptions
    { fExpression :: String
    }

data DumpOutputFormat
  = DumpPretty
  | DumpJSON
  | DumpYAML
  | DumpShow
  deriving (Bounded, Enum, Read, Show)

doDump :: DumpOptions -> IO ()
doDump DumpOptions {..} =
  runConduitRes $
  CB.sourceHandle stdin .| conduitGet (get :: Get Pkt) .|
  case outputformat of
    DumpPretty -> prettyPrinter
    DumpJSON -> jsonSink
    DumpYAML -> yamlSink
    DumpShow -> printer

-- Print every input value to standard output.
printer :: (Show a, MonadIO m) => ConduitM a Void m ()
printer = CL.mapM_ (liftIO . print)

prettyPrinter :: (Pretty a, MonadIO m) => ConduitM a Void m ()
prettyPrinter =
  CL.mapM_ (liftIO . hPutDoc stdout . (<> hardline) . group . pretty)

jsonSink :: (A.ToJSON a, MonadIO m) => ConduitM a Void m ()
jsonSink = CL.mapM_ (liftIO . BL.putStr . flip BL.snoc 0x0a . A.encode)

yamlSink :: (Y.ToJSON a, MonadIO m) => ConduitM a Void m ()
yamlSink = CL.mapM_ (liftIO . B.putStr . flip B.snoc 0x0a . Y.encode)

doFilter :: FilteringOptions -> IO ()
doFilter fo =
  runConduitRes $
  CB.sourceHandle stdin .| conduitGet (get :: Get Pkt) .|
  conduitPktFilter (parseExpressions fo) .|
  CL.map put .|
  conduitPut .|
  CB.sinkHandle stdout

doP :: Parser DumpOptions
doP =
  DumpOptions <$>
  option
    (prependAuto "Dump")
    (long "output-format" <>
     metavar "FORMAT" <>
     value DumpPretty <> showDefaultWith (drop 4 . show) <> ofHelp)
  where
    ofHelp =
      helpDoc . Just . toAnsiWlPprint $
      pretty "output format" <>
      hardline <> list (map (pretty . drop 4 . show) ofchoices)
    ofchoices = [minBound .. maxBound] :: [DumpOutputFormat]

foP :: Parser FilteringOptions
foP =
  FilteringOptions <$> argument str (metavar "EXPRESSION" <> filterTargetHelp)
  where
    filterTargetHelp =
      helpDoc . Just . toAnsiWlPprint $
      pretty "packet filter expression" <+>
      softline <> pretty "see source for current syntax"

dispatch :: Command -> IO ()
dispatch c = (banner' stderr >> hFlush stderr) >> dispatch' c
  where
    dispatch' (DumpC o) = doDump o
    dispatch' DeArmorC = doDeArmor
    dispatch' (ArmorC o) = doArmor o
    dispatch' (FilterC o) = doFilter o

main :: IO ()
main = do
  hSetBuffering stderr LineBuffering
  customExecParser
    (prefs showHelpOnError)
    (info
       (helper <*> versioner "hot" <*> cmd)
       (headerDoc (Just (toAnsiWlPprint (banner "hot"))) <>
        progDesc "hOpenPGP OpenPGP-message Tool" <>
        footerDoc (Just (toAnsiWlPprint (warranty "hot"))))) >>=
    dispatch

cmd :: Parser Command
cmd =
  hsubparser
    (command "armor" (info (ArmorC <$> aoP) (progDesc "Armor stdin to stdout")) <>
     command
       "dearmor"
       (info (pure DeArmorC) (progDesc "Dearmor stdin to stdout")) <>
     command
       "dump"
       (info (DumpC <$> doP) (progDesc "Dump OpenPGP packets from stdin")) <>
     command
       "filter"
       (info
          (FilterC <$> foP)
          (progDesc "Filter some packets from stdin to stdout")))

banner' :: Handle -> IO ()
banner' h = hPutDoc h (banner "hot" <> hardline <> warranty "hot" <> hardline)

parseExpressions :: FilteringOptions -> FilterPredicates r a
parseExpressions FilteringOptions {..} = RPFilterPredicate (parseE fExpression)
  where
    parseE e = either (error . ("filter parse error: " ++)) id (parsePExp e)

armorTypes :: [(String, ArmorType)]
armorTypes =
  [ ("message", ArmorMessage)
  , ("pubkeyblock", ArmorPublicKeyBlock)
  , ("privkeyblock", ArmorPrivateKeyBlock)
  , ("signature", ArmorSignature)
  ]

armorTypeReader :: String -> Either String ArmorType
armorTypeReader = note "unknown armor type" . flip lookup armorTypes

aoP :: Parser ArmoringOptions
aoP =
  ArmoringOptions <$>
  optional
    (strOption
       (long "comment" <> metavar "COMMENT" <> help "ASCII armor Comment field")) <*>
  option
    (eitherReader armorTypeReader)
    (long "armor-type" <> metavar "ARMORTYPE" <> armortypeHelp)
  where
    armortypeHelp =
      helpDoc . Just . toAnsiWlPprint $
      pretty "ASCII armor type" <>
      softline <> list (map (pretty . fst) armorTypes)

data ArmoringOptions =
  ArmoringOptions
    { comment :: Maybe String
    , armortype :: ArmorType
    }

doArmor :: ArmoringOptions -> IO ()
doArmor ArmoringOptions {..} = do
  m <- runConduitRes $ CB.sourceHandle stdin .| CL.consume
  let a =
        Armor
          armortype
          (maybe [] (\x -> [("Comment", x)]) comment)
          (BL.fromChunks m)
  BL.putStr $ AA.encodeLazy [a]
