namespace CryptoFile

open System
open System.IO

/// To read and write metadata.
type internal Metadata (d: (string * string) list) =
    /// Read metadata from a stream.
    new (stream : Stream) =
        let reader = new BinaryReader (stream)
        let length = reader.ReadByte () |> int
        Metadata ([for _ in 1 .. length -> reader.ReadString (), reader.ReadString ()])

    /// Write metadata to a stream.
    member this.WriteTo (stream : Stream) =
        let writer = new BinaryWriter (stream)
        writer.Write (byte d.Length)
        List.iter (fun (key : string, value : string) -> writer.Write key; writer.Write value) d

    /// May raise 'KeyNotFoundException'.
    member this.get (key: string) : string =
        List.pick (function
                    | (k, v) when k = key -> Some (v)
                    | _ -> None) d

module API =
    module internal MetadataKeys =
        let filename = "filename"
        let modificationTime = "file-modification-time"

    let internal (@@) a1 a2 = Array.append a1 a2

    let hmacSize = 256 / 8 // [byte].
    let signatureSize = Crypto.rsaKeySize / 8 // [byte].
    let keysSize = Crypto.rsaKeySize / 8 // [byte].

    let generatKeysPair () : Key * Key = Crypto.generateRSAKeysPair ()

    /// Format of the container:
    ///  <mac><signature><encrypted keys><cyphertext>
    /// Where the sizes of the three first parts are given by 'hmacSize', 'signatureSize' and 'keysSize'.
    let encryptFile (inputFilePath : string) (outputFilePath : string) (signaturePrivKey: Key) (cryptPubKey : Key)  =
        let keyAES, keyMAC, iv = Crypto.rand 16, Crypto.rand 32, Crypto.rand 16
        let fileInfo = FileInfo (inputFilePath)
        use inputStream = fileInfo.OpenRead ()
        use outputStream = new FileStream (outputFilePath, FileMode.Create, FileAccess.Write)
        use writer = new BinaryWriter (outputStream)

        outputStream.Position <- (int64 <| hmacSize + signatureSize) // Skips mac and signature. They will be written later.

        Crypto.encryptRSA cryptPubKey (keyAES @@ keyMAC @@ iv) |> writer.Write

        // Plaintext -> cryptoStream -> hmacStream -> cyphertext.
        let hmacStream, hmac = Crypto.HMACStream keyMAC outputStream
        use cryptoStream = Crypto.encryptAES keyAES iv hmacStream
        use cryptoWriter = new BinaryWriter (cryptoStream)

        // Write the file metadata.
        let metaData = Metadata ([MetadataKeys.filename, fileInfo.Name
                                  MetadataKeys.modificationTime, fileInfo.LastWriteTimeUtc.Ticks.ToString ()])
        metaData.WriteTo cryptoStream

        // Write the content of the file.
        inputStream.CopyTo cryptoStream
        cryptoStream.FlushFinalBlock ()

        // Write the HMAC at the begining of the file.
        outputStream.Position <- 0L
        writer.Write hmac.Hash

        // Write the signature.
        Crypto.signRSA signaturePrivKey hmac.Hash |> writer.Write

    /// May raise one of the following error:
    ///  * IntegrityError
    ///  * SignatureMismatch 
    ///  * UnableToDecryptAESKeys
    let decryptFile (sourceFilePath : string) (targetDirPath : string) (signaturePubKey: Key) (decryptPrivKey : Key) =
        use inputStream = new FileStream (sourceFilePath, FileMode.Open, FileAccess.Read)
        use reader = new BinaryReader (inputStream)
        let mac = reader.ReadBytes hmacSize
        let signature = reader.ReadBytes signatureSize
        let keys = 
            try reader.ReadBytes keysSize |> Crypto.decryptRSA decryptPrivKey
            with 
                | :? Security.Cryptography.CryptographicException -> raise UnableToDecryptKeys
        let keyAES = keys.[0 .. 15]
        let keyMAC = keys.[16 .. 47]
        let iv = keys.[48 .. 63]

        // Integrity validation.
        let mac' = Crypto.ComputeHMAC keyMAC inputStream
        if mac' <> mac then
            raise IntegrityError
        
        // Authentication validation.
        if not <| Crypto.verifySignRSA signaturePubKey mac' signature then
            raise SignatureMismatch

        // Decrypt metadata.
        inputStream.Position <- (int64 <| hmacSize + signatureSize + keysSize)
        use cryptoStream = Crypto.decryptAES keyAES iv inputStream
        let metadata = Metadata cryptoStream

        // Create the file and write its content and metadata.
        let filePath = Path.Combine (targetDirPath, metadata.get MetadataKeys.filename)
        let modificationTime = DateTime (metadata.get MetadataKeys.modificationTime |> int64)
        let fileInfo = FileInfo filePath
        using (fileInfo.Create ()) cryptoStream.CopyTo // We have to close the result file before updating the modification time.
        fileInfo.LastWriteTimeUtc <- modificationTime