{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE FlexibleContexts #-}
module Arbor.Albedo.Scan.Metadata
where

import Antiope.Core                          (toText, runAwsTyped)
import Antiope.Env                           (Env)
import Antiope.S3                            (S3Uri)
import Antiope.S3.Strict                     (downloadFromS3Uri)
import Control.Applicative                   ((<|>))
import Control.Lens                          (from, (&), (<&>), (^.))
import Control.Monad.Reader                  (MonadReader)
import Control.Monad.IO.Unlift               (MonadUnliftIO)
import Data.Aeson
import Data.Aeson.Types
import Data.ByteString.Lazy                  (fromStrict)
import Data.Generics.Product.Any             (the)
import Data.Generics.Product.Typed           (HasType)
import Data.Int                              (Int64)
import Data.Maybe
import Data.Thyme.Time.Core                  (UTCTime, microseconds, posixSecondsToUTCTime, toSeconds, utcTimeToPOSIXSeconds)
import GHC.Generics                          hiding (from)
import HaskellWorks.Data.Aeson
import HaskellWorks.Data.Network.Ip.Ipv4     (IpBlock)
import HaskellWorks.Data.Network.Ip.Validity (Canonical)

import qualified Data.Text as Text


newtype Proto = Proto { unProto :: Int } deriving (Generic, Show, Eq, FromJSON, ToJSON)

data Command = Command
  { name        :: Maybe String
  , scanner     :: Maybe String
  , probe       :: Maybe String
  , probeBanner :: Maybe String -- use if zmap and zgrab use different probes(ie udp and cldap), banner-commander would prioritize probeBanner
  , probeArgs   :: Maybe String
  , port        :: Int
  , proto       :: Proto
  , destination :: Either FilePath S3Uri
  , cidr        :: IpBlock Canonical
  , deadline    :: UTCTime
  , banners     :: Bool
  } deriving (Generic, Show, Eq)

data ScanMetaData = ScanMetaData
  { startAt    :: UTCTime
  , stopAt     :: UTCTime
  , provenance :: Maybe String
  , scanCmd    :: Command
  , digest     :: String
  , port       :: Int
  , proto      :: Proto
  , cidr       :: IpBlock Canonical
  , source     :: Maybe String
  , data'      :: String
  } deriving (Generic, Show, Eq)

instance FromJSON Command where
  parseJSON = withObject "object" $ \o ->
    Command <$> o .:? "name"
            <*> o .: "scanner"
            <*> o .: "probe"
            <*> o .:? "probeBanner"
            <*> o .:? "probeArgs"
            <*> o .: "port"
            <*> (o .: "proto")
            <*> (o .:  "destination" >>= readDestination)
            <*> (o .:  "cidr"        >>= readJson "cidr")
            <*> (o .:  "deadline"    & fmap int64ToUTCTime)
            <*> (o .:? "banners" .!= False)

instance ToJSON Command where
  toJSON r = objectWithoutNulls
    [ "name"        .= (r ^. the @"name")
    , "scanner"     .= (r ^. the @"scanner")
    , "probe"       .= (r ^. the @"probe")
    , "probeBanner" .= (r ^. the @"probeBanner")
    , "probeArgs"   .= (r ^. the @"probeArgs")
    , "port"        .= (r ^. the @"port")
    , "proto"       .= (r ^. the @"proto")
    , "destination" .= either Text.pack toText (r ^. the @"destination")
    , "cidr"        .= show (r ^. the @"cidr")
    , "deadline"    .= utcTimeToInt64 (r ^. the @"deadline")
    , "banners"     .= (r ^. the @"banners")
    ]

instance FromJSON ScanMetaData where
  parseJSON = withObject "object" $ \o ->
    ScanMetaData <$> (o .: "startAt" & fmap int64ToUTCTime)
                 <*> (o .: "stopAt"  & fmap int64ToUTCTime)
                 <*> (o .:? "provenance")
                 <*> o .: "scanCmd"
                 <*> o .: "digest"
                 <*> o .: "port"
                 <*> (o .: "proto")
                 <*> (o .: "cidr"   >>= readJson "cidr")
                 <*> (o .:? "source")
                 <*> (o .: "data")

instance ToJSON ScanMetaData where
  toJSON r = objectWithoutNulls
    [ "startAt"    .= utcTimeToInt64 (r ^. the @"startAt")
    , "stopAt"     .= utcTimeToInt64 (r ^. the @"stopAt")
    , "provenance" .= (r ^. the @"provenance")
    , "scanCmd"    .= (r ^. the @"scanCmd")
    , "digest"     .= (r ^. the @"digest")
    , "port"       .= (r ^. the @"port")
    , "proto"      .= (r ^. the @"proto")
    , "cidr"       .= show (r ^. the @"cidr")
    , "source"     .= (r ^. the @"source")
    , "data"       .= (r ^. the @"data'")
    ]

readDestination :: String -> Parser (Either FilePath S3Uri)
readDestination s = (Right <$> readJson "destination (s3uri)"    s)
                <|> (Left  <$> readJson "destination (filepath)" s)

int64ToUTCTime :: Int64 -> UTCTime
int64ToUTCTime a = posixSecondsToUTCTime $ (a * 1000 * 1000) ^. from microseconds

utcTimeToInt64 :: UTCTime -> Int64
utcTimeToInt64 ut = round (toSeconds (utcTimeToPOSIXSeconds ut) :: Double)

loadScanMetadata :: (MonadUnliftIO m, MonadReader r m, HasType Env r)
    => S3Uri
    -> m (Maybe ScanMetaData)
loadScanMetadata uri = runAwsTyped $ do
    mb <- downloadFromS3Uri uri
    pure $ mb >>= (decode . fromStrict)