//
//  CipherChaCha20Tests.swift
//  CryptoSwift
//
//  Created by Marcin Krzyzanowski on 27/12/14.
//  Copyright (c) 2014 Marcin Krzyzanowski. All rights reserved.
//
import XCTest
@testable import CryptoSwift

final class ChaCha20Tests: XCTestCase {

    override func setUp() {
        super.setUp()
    }
    
    override func tearDown() {
        super.tearDown()
    }

    func testChaCha20() {
        let keys:[Array<UInt8>] = [
            [0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00],
            [0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x01],
            [0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00],
            [0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00],
            [0x00,0x01,0x02,0x03,0x04,0x05,0x06,0x07,0x08,0x09,0x0A,0x0B,0x0C,0x0D,0x0E,0x0F,0x10,0x11,0x12,0x13,0x14,0x15,0x16,0x17,0x18,0x19,0x1A,0x1B,0x1C,0x1D,0x1E,0x1F]
        ]
        
        let ivs:[Array<UInt8>] = [
            [0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00],
            [0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00],
            [0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x01],
            [0x01,0x00,0x00,0x00,0x00,0x00,0x00,0x00],
            [0x00,0x01,0x02,0x03,0x04,0x05,0x06,0x07]
        ]
        
        let expectedHexes = [
            "76B8E0ADA0F13D90405D6AE55386BD28BDD219B8A08DED1AA836EFCC8B770DC7DA41597C5157488D7724E03FB8D84A376A43B8F41518A11CC387B669B2EE6586",
            "4540F05A9F1FB296D7736E7B208E3C96EB4FE1834688D2604F450952ED432D41BBE2A0B6EA7566D2A5D1E7E20D42AF2C53D792B1C43FEA817E9AD275AE546963",
            "DE9CBA7BF3D69EF5E786DC63973F653A0B49E015ADBFF7134FCB7DF137821031E85A050278A7084527214F73EFC7FA5B5277062EB7A0433E445F41E3",
            "EF3FDFD6C61578FBF5CF35BD3DD33B8009631634D21E42AC33960BD138E50D32111E4CAF237EE53CA8AD6426194A88545DDC497A0B466E7D6BBDB0041B2F586B",
            "F798A189F195E66982105FFB640BB7757F579DA31602FC93EC01AC56F85AC3C134A4547B733B46413042C9440049176905D3BE59EA1C53F15916155C2BE8241A38008B9A26BC35941E2444177C8ADE6689DE95264986D95889FB60E84629C9BD9A5ACB1CC118BE563EB9B3A4A472F82E09A7E778492B562EF7130E88DFE031C79DB9D4F7C7A899151B9A475032B63FC385245FE054E3DD5A97A5F576FE064025D3CE042C566AB2C507B138DB853E3D6959660996546CC9C4A6EAFDC777C040D70EAF46F76DAD3979E5C5360C3317166A1C894C94A371876A94DF7628FE4EAAF2CCB27D5AAAE0AD7AD0F9D4B6AD3B54098746D4524D38407A6DEB3AB78FAB78C9"
        ]

        for idx in 0..<keys.count {            
            let expectedHex = expectedHexes[idx]
            let message = Array<UInt8>(repeating: 0, count: (expectedHex.characters.count / 2))
            
            do {
                let encrypted = try ChaCha20(key: keys[idx], iv: ivs[idx]).encrypt(message)
                let decrypted = try ChaCha20(key: keys[idx], iv: ivs[idx]).decrypt(encrypted)
                XCTAssertEqual(message, decrypted, "ChaCha20 decryption failed");
                
                // check extension
                let messageData = Data(bytes: message);
                let encrypted2 = try messageData.encrypt(cipher: ChaCha20(key: keys[idx], iv: ivs[idx]))
                XCTAssertNotNil(encrypted2, "")
                XCTAssertEqual(Data(bytes: encrypted), encrypted2, "ChaCha20 extension failed")
            } catch CipherError.encrypt {
                XCTAssert(false, "Encryption failed")
            } catch CipherError.decrypt {
                XCTAssert(false, "Decryption failed")
            } catch {
                XCTAssert(false, "Failed")
            }
        }
    }

    func testVector1Py() {
        let key:Array<UInt8> = [0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00]
        let iv:Array<UInt8>  = [0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00]
        let expected:Array<UInt8> = [0x76,0xB8,0xE0,0xAD,0xA0,0xF1,0x3D,0x90,0x40,0x5D,0x6A,0xE5,0x53,0x86,0xBD,0x28,0xBD,0xD2,0x19,0xB8,0xA0,0x8D,0xED,0x1A,0xA8,0x36,0xEF,0xCC,0x8B,0x77,0x0D,0xC7,0xDA,0x41,0x59,0x7C,0x51,0x57,0x48,0x8D,0x77,0x24,0xE0,0x3F,0xB8,0xD8,0x4A,0x37,0x6A,0x43,0xB8,0xF4,0x15,0x18,0xA1,0x1C,0xC3,0x87,0xB6,0x69,0xB2,0xEE,0x65,0x86]
        let message = Array<UInt8>(repeating: 0, count: expected.count)

        do {
            let encrypted = try ChaCha20(key: key, iv: iv).encrypt(message)
            XCTAssertEqual(encrypted, expected, "Ciphertext failed")
        } catch {
            XCTFail()
        }
    }

    func testChaCha20_encrypt_partial() {
        let key:Array<UInt8> = [0x2b,0x7e,0x15,0x16,0x28,0xae,0xd2,0xa6,0xab,0xf7,0x15,0x88,0x09,0xcf,0x4f,0x3c];
        let iv:Array<UInt8> = [0x00,0x01,0x02,0x03,0x04,0x05,0x06,0x07,0x08,0x09,0x0A,0x0B,0x0C,0x0D,0x0E,0x0F]
        let plaintext:Array<UInt8> = [0x6b,0xc1,0xbe,0xe2,0x2e,0x40,0x9f,0x96,0xe9,0x3d,0x7e,0x11,0x73,0x93,0x17,0x2a,0x6b,0xc1,0xbe,0xe2,0x2e,0x40,0x9f,0x96,0xe9,0x3d,0x7e,0x11,0x73,0x93,0x17,0x2a]

        do {
            let cipher = try ChaCha20(key: key, iv:iv)

            var ciphertext = Array<UInt8>()
            var encryptor = cipher.makeEncryptor()
            ciphertext += try encryptor.update(withBytes: Array(plaintext[0..<8]))
            ciphertext += try encryptor.update(withBytes: Array(plaintext[8..<16]))
            ciphertext += try encryptor.update(withBytes: Array(plaintext[16..<32]))
            ciphertext += try encryptor.finish()
            XCTAssertEqual(try cipher.encrypt(plaintext), ciphertext, "encryption failed")
        } catch {
            XCTFail()
        }
    }

    func testChaCha20Performance() {
        let key:Array<UInt8> = [0x2b,0x7e,0x15,0x16,0x28,0xae,0xd2,0xa6,0xab,0xf7,0x15,0x88,0x09,0xcf,0x4f,0x3c];
        let iv:Array<UInt8> = [0x00,0x01,0x02,0x03,0x04,0x05,0x06,0x07,0x08,0x09,0x0A,0x0B,0x0C,0x0D,0x0E,0x0F]
        let message = Array<UInt8>(repeating: 7, count: (1024 * 1024) * 1)
        measureMetrics([XCTPerformanceMetric_WallClockTime], automaticallyStartMeasuring: true, for: { () -> Void in
            do {
                let _ = try ChaCha20(key: key, iv: iv).encrypt(message)
            } catch {
                XCTFail()
            }
            self.stopMeasuring()
        })
    }
}