{-# LANGUAGE OverloadedStrings #-}

-- Application runs as a CGI script and enables saving an
-- encrypted text to a sqlite3 database.
module Main where

import Control.Monad (replicateM)
import Crypto.PBKDF.ByteString (sha256PBKDF2)
import System.Random (randomIO)

import qualified Codec.Crypto.AES as AES
import qualified Data.Aeson as JSON
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as LBS
import qualified Data.ByteString.Char8 as Char8
import qualified Database.SQLite.Simple as SQLite
import qualified Database.SQLite.Simple.FromRow as SQLite
import qualified Network.CGI as CGI



-- * Types and instances


newtype Password = Password BS.ByteString
    deriving (Show)


-- Hashed crypto key
newtype Key = Key BS.ByteString
    deriving (Show)


-- DB primary key
newtype PK = PK Int
    deriving (Show)


-- Password salt
newtype Salt = Salt BS.ByteString
    deriving (Show)


-- Initialization vector
newtype IV = IV BS.ByteString
    deriving (Show)


newtype Plaintext = Plaintext BS.ByteString
    deriving (Show)


newtype Ciphertext = Ciphertext BS.ByteString
    deriving (Show)


-- A database row
data Text = Text PK Salt IV Ciphertext
    deriving (Show)


data Request
    = RequestDisplay Password -- Display a text without saving
    | RequestSave Password Plaintext -- Save a text
    | RequestHelp -- In all other cases, interpret as a help request
    deriving (Show)


data Response
    = ResponsePlaintext Plaintext -- Sent for successful requests
    | ResponseHelp BS.ByteString -- Sent in other cases, including false passwords
    deriving (Show)


data DecryptResult
    = DecryptSuccessful Plaintext
    | DecryptNoTextsSaved
    | DecryptWrongKey
    deriving (Show)


instance SQLite.FromRow Text where
    fromRow = Text
        <$> (PK         <$> SQLite.field)
        <*> (Salt       <$> SQLite.field)
        <*> (IV         <$> SQLite.field)
        <*> (Ciphertext <$> SQLite.field)


instance JSON.ToJSON Response where
    toJSON (ResponsePlaintext (Plaintext pt)) = JSON.object ["text" JSON..= Char8.unpack pt]
    toJSON (ResponseHelp text)                = JSON.object ["help" JSON..= Char8.unpack text]



-- * Main program


main :: IO ()
main = do
    dbConnection <- SQLite.open "db/db.sqlite3"

    CGI.runCGI $ do
        CGI.setHeader "Content-type" "application/json"

        request  <- parseRequest 
        response <- CGI.liftIO $ handleRequest dbConnection request

        CGI.outputFPS $ JSON.encode response

    return ()

    where
        -- Parse any password and any text sent to construct a Request
        parseRequest :: CGI.CGI Request
        parseRequest = do
            maybePassword  <- fmap (Password  . LBS.toStrict) <$> CGI.getInputFPS "password"
            maybePlaintext <- fmap (Plaintext . LBS.toStrict) <$> CGI.getInputFPS "save"

            return $ case maybePassword of
                Just password -> case maybePlaintext of
                    Just plaintext -> RequestSave password plaintext
                    Nothing        -> RequestDisplay password
                Nothing  -> RequestHelp


-- Request handler

handleRequest :: SQLite.Connection -> Request -> IO Response

-- Handle help request (request when no key was given)
handleRequest dbConnection RequestHelp = return helpResponse

-- Handle save request (try to save the text with the key.
-- If it works, save the new text and display the submitted text
handleRequest dbConnection (RequestSave password plaintext) = do
    decryptResult <- attemptDecrypt dbConnection password

    case decryptResult of
        DecryptWrongKey -> return helpResponse

        DecryptNoTextsSaved -> do
            (Salt s, IV iv, Ciphertext ct) <- encryptData password plaintext

            SQLite.execute dbConnection
                "INSERT INTO texts (salt, iv, text) VALUES (?, ?, ?)"
                (s, iv, ct)

            return $ ResponsePlaintext plaintext

        DecryptSuccessful _ -> do
            (Salt s, IV iv, Ciphertext ct) <- encryptData password plaintext

            SQLite.execute dbConnection
                "UPDATE texts SET salt = ?, iv = ?, text = ? WHERE id = ?"
                (s, iv, ct, 1 :: Int)

            return $ ResponsePlaintext plaintext

    where
        encryptData :: Password -> Plaintext -> IO (Salt, IV, Ciphertext)
        encryptData password plaintext = do
            -- Generate salt
            salt <- Salt . BS.pack <$> replicateM 64 randomIO

            (initializationVector, ciphertext) <- encryptWithAES
                (hashPassword salt password)
                plaintext

            return (salt, initializationVector, ciphertext)
            

-- Handle display request (display text)
handleRequest dbConnection (RequestDisplay password) = do
    decryptResult <- attemptDecrypt dbConnection password

    return $ case decryptResult of
        DecryptWrongKey             -> helpResponse
        DecryptNoTextsSaved         -> ResponsePlaintext $ Plaintext ""
        DecryptSuccessful plaintext -> ResponsePlaintext plaintext


-- Response helper functions

helpResponse = ResponseHelp "Usage: POST ?password=PASSWORD&save=TEXT. This message is also returned for invalid passwords."


-- Database/auth helper functions

attemptDecrypt :: SQLite.Connection -> Password -> IO DecryptResult
attemptDecrypt dbConnection password = do
    texts <- SQLite.query_ dbConnection "SELECT * FROM texts" :: IO [Text]

    case texts of
        ((Text _ salt initializationVector ciphertext):_) -> do
            let key = hashPassword salt password
            let maybePlaintext = decryptWithAES key initializationVector ciphertext

            case maybePlaintext of
                Just plaintext -> return $ DecryptSuccessful plaintext
                Nothing        -> return $ DecryptWrongKey

        _ -> return DecryptNoTextsSaved


-- * Crypto helper functions


hashPassword :: Salt -> Password -> Key
hashPassword (Salt s) (Password pw) = Key $ sha256PBKDF2 pw s 100000 32


-- Extra string prepended during encryption to make easy later
-- checking of whether the decryption was successful possible.
extraEncryptString = "---SUCCESSFUL-DECRYPT---"


decryptWithAES
    :: Key
    -> IV
    -> Ciphertext
    -> Maybe Plaintext
decryptWithAES (Key k) (IV iv) (Ciphertext ct) =
    let pt = AES.crypt' AES.CTR k iv AES.Decrypt ct
        extraStringLength  = BS.length extraEncryptString

    in if BS.take extraStringLength pt == extraEncryptString
        then Just $ Plaintext $ BS.drop extraStringLength pt
        else Nothing


encryptWithAES
    :: Key
    -> Plaintext
    -> IO (IV, Ciphertext)
encryptWithAES (Key k) (Plaintext pt) = do
    initializationVector@(IV iv) <- generateInitializationVector

    let ciphertext = Ciphertext $ AES.crypt' AES.CTR k iv AES.Encrypt $
            BS.append extraEncryptString pt

    return (initializationVector, ciphertext)

    where
        generateInitializationVector :: IO IV
        generateInitializationVector =
            IV . BS.pack <$> replicateM 16 randomIO

