-----------------------------------------------------------------------------
-- |
-- Module      :  Network.AWS.Authentication
-- Copyright   :  (c) Greg Heartsfield 2007, David Himmelstrup 2009
-- License     :  BSD3
--
-- Implements authentication and low-level communication with Amazon
-- Web Services, such as S3, SimpleDB, EC2, and others.
-- API Version 2009-04-15
-- <http://docs.amazonwebservices.com/AmazonSimpleDB/2009-04-15/DeveloperGuide/>
-----------------------------------------------------------------------------

module Network.AWS.Authentication (
   -- * Function Types
   runAction, --isAmzHeader, -- preSignedURI,
   -- * Data Types
   SimpleDBAction(..),
   -- Misc functions
   -- mimeEncodeQP, mimeDecode
   ) where

import Network.AWS.AWSResult
import Network.AWS.AWSConnection
import Network.AWS.ArrowUtils
import Network.HTTP as HTTP hiding (simpleHTTP_)
import Network.HTTP.HandleStream (simpleHTTP_)
import Network.Stream (Result)
import Network.URI as URI
import qualified Data.ByteString.Lazy.Char8 as L

import Data.ByteString.Char8 (pack, unpack)

import Data.HMAC
import Codec.Binary.Base64 (encode, decode)
import Codec.Utils (Octet)

import Data.Char (intToDigit, digitToInt, ord, chr, toLower)
import Data.Bits ((.&.))
import qualified Codec.Binary.UTF8.String as US

import Data.List (sortBy, groupBy, intersperse, intercalate, sort)
import Data.Maybe

import System.Time
import System.Locale

import Text.XML.HXT.Arrow

-- | An action to be performed using SimpleDB
data SimpleDBAction =
    SimpleDBAction {
      -- | Connection and authentication information
      sdbConnection :: AWSConnection,
      -- | Query parameters (requires a prefix of @?@)
      sdbQuery :: [String],
      -- | Additional header fields to send
      sdbMetaData :: [(String, String)],
      -- | Body of action, if sending data
      sdbBody :: L.ByteString,
      -- | Type of action, 'PUT', 'GET', etc.
      sdbOperation :: RequestMethod
    } deriving (Show)

-- | Transform an 'SimpleDBAction' into an HTTP request.  Does not add
--   authentication or date information, so it is not suitable for
--   sending directly to AWS.
requestFromAction :: SimpleDBAction -- ^ Action to transform
                  -> HTTP.HTTPRequest L.ByteString -- ^ Action represented as an HTTP Request.
requestFromAction a =
    Request { rqURI = URI { uriScheme = "",
                            uriAuthority = Nothing,
                            uriPath = qpath,
                            uriQuery = '?':intercalate "&" (sdbQuery a),
                            uriFragment = "" },
              rqMethod = sdbOperation a,
              rqHeaders = Header HdrHost (sdbHostname a) :
                          headersFromAction a,
              rqBody = (sdbBody a)
            }
    where qpath = ['/']

-- | Create 'Header' objects from an action.
headersFromAction :: SimpleDBAction
                  -> [Header]
headersFromAction = map (\(k,v) -> case k of
                                    "Content-Type" -> Header HdrContentType v
                                    otherwise -> Header (HdrCustom k) (mimeEncodeQP v))
                    . sdbMetaData

-- | Inspect HTTP body, and add a @Content-Length@ header with the
--   correct length.
addContentLengthHeader :: HTTP.HTTPRequest L.ByteString -> HTTP.HTTPRequest L.ByteString
addContentLengthHeader req = insertHeader HdrContentLength (show (L.length (rqBody req))) req

-- | Add AWS authentication header to an HTTP request.
addAuthenticationHeader :: SimpleDBAction     -- ^ Action with authentication data
                        -> HTTP.HTTPRequest L.ByteString -- ^ Request to transform
                        -> HTTP.HTTPRequest L.ByteString -- ^ Authenticated request
addAuthenticationHeader act req
    = appendQuery ("&Signature="++ urlEncode signature) $
      req -- insertHeader HdrAuthorization auth_string req
    where auth_string = "AWS " ++ awsAccessKey conn ++ ":" ++ signature
          signature = makeSignature conn (stringToSign act req)
          conn = sdbConnection act

appendQuery :: String -> HTTP.HTTPRequest L.ByteString -> HTTP.HTTPRequest L.ByteString
appendQuery string req
    = req{ rqURI = uri{ uriQuery = uriQuery uri ++ string } }
    where uri = rqURI req

-- | Sign a string using the given authentication data
makeSignature :: AWSConnection -- ^ Action with authentication data
              -> String -- ^ String to sign
              -> String -- ^ Base-64 encoded signature
makeSignature c s =
        encode (hmac_sha1 keyOctets msgOctets)
        where keyOctets = string2words (awsSecretKey c)
              msgOctets = string2words s

-- | Generate text that will be signed and subsequently added to the
--   request.
stringToSign :: SimpleDBAction -> HTTP.HTTPRequest L.ByteString -> String
stringToSign a r =
    canonicalizeHeaders r ++
    canonicalizeResource a ++ "\n" ++
    intercalate "&" (sort (sdbQuery a))

-- | Extract header data needed for signing.
canonicalizeHeaders :: HTTP.HTTPRequest L.ByteString -> String
canonicalizeHeaders r =
    http_verb ++ "\n" ++
    get_header HdrHost ++ "\n"
--    hdr_content_md5 ++ "\n" ++
--    hdr_content_type ++ "\n" ++
--    dateOrExpiration ++ "\n"
        where http_verb = show (rqMethod r)
              hdr_content_md5 = get_header HdrContentMD5
              hdr_date = get_header HdrDate
              hdr_content_type = get_header HdrContentType
              get_header h = fromMaybe "" (findHeader h r)
              dateOrExpiration = fromMaybe hdr_date (findHeader HdrExpires r)

-- | Combine same-named headers.
combineHeaders :: [[(String, String)]] -> [(String, String)]
combineHeaders = map mergeSameHeaders

-- | Headers with same name should have values merged.
mergeSameHeaders :: [(String, String)] -> (String, String)
mergeSameHeaders h@(x:_) = let values = map snd h
                     in ((fst x), (concat $ intersperse "," values))

-- | Group headers with the same name.
groupHeaders :: [(String, String)] -> [[(String, String)]]
groupHeaders = groupBy (\a b -> fst a == fst b)

-- | Sort by key name.
sortHeaders :: [(String, String)] -> [(String, String)]
sortHeaders = sortBy (\a b -> fst a `compare` fst b)

-- | Make 'Header' easier to work with, and lowercase keys.
headerToLCKeyValue :: Header -> (String, String)
headerToLCKeyValue (Header k v) = (map toLower (show k), v)

-- | Determine if a header belongs in the StringToSign
isAmzHeader :: Header -> Bool
isAmzHeader h =
    case h of
      Header (HdrCustom k) _ -> isPrefix amzHeader k
      otherwise -> False

-- | is the first list a prefix of the second?
isPrefix :: Eq a => [a] -> [a] -> Bool
isPrefix a b = a == take (length a) b

-- | Prefix used by Amazon metadata headers
amzHeader :: String
amzHeader = "x-amz-"

-- | Extract resource name, as required for signing.
canonicalizeResource :: SimpleDBAction -> String
canonicalizeResource a = uri
    where uri = '/' : []


-- | Add a date string to a request.
addDateToReq :: HTTP.HTTPRequest L.ByteString -- ^ Request to modify
             -> String       -- ^ Date string, in RFC 2616 format
             -> String       -- ^ Timestamp
             -> HTTP.HTTPRequest L.ByteString-- ^ Request with date header added
addDateToReq r d ts
    = -- appendQuery ("&Timestamp=" ++ urlEncode ts) $
      r {HTTP.rqHeaders =
         HTTP.Header HTTP.HdrDate d : HTTP.rqHeaders r}

-- | Add an expiration date to a request.
addExpirationToReq :: HTTP.HTTPRequest L.ByteString -> String -> HTTP.HTTPRequest L.ByteString
addExpirationToReq r = addHeaderToReq r . HTTP.Header HTTP.HdrExpires

-- | Attach an HTTP header to a request.
addHeaderToReq :: HTTP.HTTPRequest L.ByteString -> Header -> HTTP.HTTPRequest L.ByteString
addHeaderToReq r h = r {HTTP.rqHeaders = h : HTTP.rqHeaders r}

-- | Get hostname to connect to. Needed for european buckets
sdbHostname :: SimpleDBAction -> String
sdbHostname a = sdbhost
    where sdbhost = awsHost (sdbConnection a)

-- | Get current time in HTTP 1.1 format (RFC 2616)
--   Numeric time zones should be used, but I'd rather not subvert the
--   intent of ctTZName, so we stick with the name format.  Otherwise,
--   we could send @+0000@ instead of @GMT@.
--   see:
--   <http://www.ietf.org/rfc/rfc2616.txt>
--   <http://www.ietf.org/rfc/rfc1123.txt>
--   <http://www.ietf.org/rfc/rfc822.txt>
httpCurrentDate :: IO String
httpCurrentDate =
    do c <- getClockTime
       let utc_time = (toUTCTime c) {ctTZName = "GMT"}
       return $ formatCalendarTime defaultTimeLocale rfc822DateFormat utc_time

httpCurrentTimestamp :: IO String
httpCurrentTimestamp =
    do c <- getClockTime
       let utc_time = (toUTCTime c) {ctTZName = "GMT"}
       return $ formatCalendarTime defaultTimeLocale "%Y-%m-%dT%H:%M:%S" utc_time


-- | Convenience for dealing with HMAC-SHA1
string2words :: String -> [Octet]
string2words = US.encode

-- | Construct the request specified by an SimpleDBAction, send to Amazon,
--   and return the response.  Todo: add MD5 signature.
runAction :: SimpleDBAction -> IO (AWSResult (HTTPResponse L.ByteString))
runAction a = runAction' a (sdbHostname a)

runAction' :: SimpleDBAction -> String -> IO (AWSResult (HTTPResponse L.ByteString))
runAction' a hostname = do
        c <- (openTCPConnection hostname (awsPort (sdbConnection a)))
--bufferOps = lazyBufferOp
        cd <- httpCurrentDate
        ts <- httpCurrentTimestamp
        let a' = a{sdbQuery = ("Timestamp="++urlEncode ts) : ("AWSAccessKeyId="++awsAccessKey (sdbConnection a)) : sdbQuery a}
        let aReq = addAuthenticationHeader a' $
                   addContentLengthHeader $
                   addDateToReq (requestFromAction a') cd ts
        --putStrLn (stringToSign a' aReq)
        --print aReq
        result <- simpleHTTP_ c aReq
        --print result
        close c
        createAWSResult a result

-- | Construct a pre-signed URI, but don't act on it.  This is useful
--   for when an expiration date has been set, and the URI needs to be
--   passed on to a client.
preSignedURI :: SimpleDBAction -- ^ Action with resource
             -> Integer -- ^ Expiration time, in seconds since
                        --   00:00:00 UTC on January 1, 1970
             -> URI -- ^ URI of resource
preSignedURI a e =
    let c = (sdbConnection a)
        srv = (awsHost c)
        pt = (show (awsPort c))
        accessKeyQuery = "AWSAccessKeyId=" ++ awsAccessKey c
        beginQuery = case (sdbQuery a) of
                  [] -> "?"
                  x -> intercalate "&" x ++ "&"
        expireQuery = "Expires=" ++ show e
        toSign = "GET\n\n\n" ++ show e ++ "\n/"
        sigQuery = "Signature=" ++ urlEncode (makeSignature c toSign)
        q = beginQuery ++ accessKeyQuery ++ "&" ++
                expireQuery ++ "&" ++ sigQuery
    in URI { uriScheme = "http:",
             uriAuthority = Just (URIAuth "" srv (':' : pt)),
             uriPath = "/",
             uriQuery = q,
             uriFragment = ""
           }


-- | Inspect a response for network errors, HTTP error codes, and
--   Amazon error messages.
--   We need the original action in case we get a 307 (temporary redirect)
createAWSResult :: SimpleDBAction -> Result (HTTPResponse L.ByteString) -> IO (AWSResult (HTTPResponse L.ByteString))
createAWSResult a b = either handleError handleSuccess b
    where handleError = return . Left . NetworkError
          handleSuccess s = case (rspCode s) of
                              (2,_,_) -> return (Right s)
                              -- temporary redirect
                              (3,0,7) -> case (findHeader HdrLocation s) of
                                                Just l -> runAction' a (getHostname l)
                                                Nothing -> return (Left $ AWSError "Temporary Redirect" "Redirect without location header")  -- not good
                              (4,0,4) -> return (Left $ AWSError "NotFound" "404 Not Found")  -- no body, so no XML to parse
                              otherwise -> do e <- parseRestErrorXML (L.unpack (rspBody s))
                                              return (Left e)
          -- Get hostname part from http url.
          getHostname :: String -> String
          getHostname h = case parseURI h of
                             Just u -> case (uriAuthority u) of
                                           Just auth -> (uriRegName auth)
                                           Nothing -> ""
                             Nothing -> ""
-- | Find the errors embedded in an XML message body from Amazon.
parseRestErrorXML :: String -> IO ReqError
parseRestErrorXML x =
    do e <- runX (readString [(a_validate,v_0)] x
                                 >>> processRestError)
       case e of
         [] -> return (AWSError "NoErrorInMsg"
                       ("HTTP Error condition, but message body"
                        ++ "did not contain error code."))
         x:xs -> return x


-- | Find children of @Error@ entity, use their @Code@ and @Message@
--   entities to create an 'AWSError'.
processRestError = deep (isElem >>> hasName "Error") >>>
                   split >>> first (text <<< atTag "Code") >>>
                   second (text <<< atTag "Message") >>>
                   unsplit (\x y -> AWSError x y)
{-
processRestError = proc n -> do deep (isElem >>> hasName "Error") -< n
                                x <- text <<< atTag "Code" -< n
                                y <- text <<< atTag "Message" -< n
                                returnA -< (AWSError x y)
-}

--- mime header encoding
mimeEncodeQP, mimeDecode :: String -> String
-- | Decode a mime string, we know about quoted printable and base64 encoded UTF-8
-- S3 may convert quoted printable to base64
mimeDecode a
    | isPrefix utf8qp a =
        mimeDecodeQP $ encoded_payload utf8qp a
    | isPrefix utf8b64 a =
        mimeDecodeB64 $ encoded_payload utf8b64 a
    | otherwise =
        a
    where
      utf8qp  = "=?UTF-8?Q?"
      utf8b64 = "=?UTF-8?B?"
      -- '=?UTF-8?Q?foobar?=' -> 'foobar'
      encoded_payload prefix = reverse . drop 2 . reverse . drop (length prefix)

mimeDecodeQP :: String -> String
mimeDecodeQP =
    US.decodeString . mimeDecodeQP'

mimeDecodeQP' :: String -> String
mimeDecodeQP' ('=':a:b:rest) =
    chr (16 * digitToInt a + digitToInt b)
            : mimeDecodeQP' rest
mimeDecodeQP' (h:t) =h : mimeDecodeQP' t
mimeDecodeQP' [] = []

mimeDecodeB64 :: String -> String
mimeDecodeB64 s =
    case decode s of
        Nothing -> ""
        Just a ->  US.decode a

 -- Encode a String into quoted printable, if needed.
 -- eq: =?UTF-8?Q?=aa?=
mimeEncodeQP s =
    if any reservedChar s
    then "=?UTF-8?Q?" ++ (mimeEncodeQP' $ US.encodeString s) ++ "?="
    else s

mimeEncodeQP' :: String -> String
mimeEncodeQP' [] = []
mimeEncodeQP' (h:t) =
    let str = if reservedChar h then escape h else [h]
    in str ++ mimeEncodeQP' t
    where
        escape x =
            let y = ord x in
            [ '=', intToDigit ((y `div` 16) .&. 0xf), intToDigit (y .&. 0xf) ]

-- Char needs escaping?
reservedChar :: Char -> Bool
reservedChar x
    -- from space (0x20) till '~' everything is fine. The rest are control chars, or high bit.
    | xi >= 0x20 && xi <= 0x7e = False
    | otherwise = True
    where xi = ord x