//

//  MYCryptoTest.m

//  MYCrypto-iPhone

//

//  Created by Jens Alfke on 4/1/09.

//  Copyright 2009 Jens Alfke. All rights reserved.

//

#import "MYPublicKey.h"

#import "MYPrivateKey.h"

#import "MYKeychain.h"

#import "MYDigest.h"

#import "MYIdentity.h"

#if !TARGET_OS_IPHONE

#import "MYCrypto+Cocoa.h"

#endif

#import "MYCrypto_Private.h"

#if DEBUG

#define kTestCaseRSAKeySize 2048

#pragma mark -

#pragma mark KEYCHAIN:

TestCase(MYKeychain) {

    MYKeychain *kc = [MYKeychain defaultKeychain];

    Log(@"Default keychain = %@", kc);

    CAssert(kc);

#if !MYCRYPTO_USE_IPHONE_API

    CAssert(kc.path);

#endif

    kc = [MYKeychain allKeychains];

    Log(@"All-keychains = %@", kc);

    CAssert(kc);

#if !MYCRYPTO_USE_IPHONE_API

    CAssertEq(kc.path,nil);

#endif

}

TestCase(EnumerateKeys) {

    RequireTestCase(MYKeychain);

    NSEnumerator *e = [[MYKeychain allKeychains] enumeratePublicKeys];

    Log(@"Public Key Enumerator = %@", e);

    CAssert(e);

    for (MYPublicKey *key in e) {

        Log(@"Found %@ -- name=%@", key, key.name);

    }

    e = [[MYKeychain allKeychains] enumeratePrivateKeys];

    Log(@"Key-Pair Enumerator = %@", e);

    CAssert(e);

    for (MYPrivateKey *key in e) {

        Log(@"Found %@ -- name=%@", key, key.name);

    }

    e = [[MYKeychain allKeychains] enumerateSymmetricKeys];

    Log(@"Symmetric Key Enumerator = %@", e);

    CAssert(e);

    for (MYSymmetricKey *key in e) {

        Log(@"Found %@ -- name=%@", key, key.name);

    }

}

TestCase(EnumerateCerts) {

    RequireTestCase(MYKeychain);

    NSEnumerator *e = [[MYKeychain allKeychains] enumerateCertificates];

    Log(@"Enumerator = %@", e);

    CAssert(e);

    for (MYCertificate *cert in e) {

        //Log(@"Found %@ -- name=%@, email=%@", cert, cert.commonName, cert.emailAddresses);

    }

}

TestCase(EnumerateIdentities) {

    RequireTestCase(MYKeychain);

    NSEnumerator *e = [[MYKeychain allKeychains] enumerateIdentities];

    Log(@"Enumerator = %@", e);

    CAssert(e);

    for (MYIdentity *ident in e) {

        Log(@"Found %@\n\tcommonName=%@\n\temails=(%@)\n\tkey=%@",

            ident, ident.commonName, 

#if TARGET_OS_IPHONE

            nil,

#else

            [ident.emailAddresses componentsJoinedByString: @", "],

#endif

            ident.privateKey);

    }

}

#pragma mark -

#pragma mark SYMMETRIC KEYS:

static void testSymmetricKey( CCAlgorithm algorithm, unsigned sizeInBits, MYKeychain *inKeychain ) {

    NSAutoreleasePool *pool = [NSAutoreleasePool new];

    MYSymmetricKey *key = nil;

    @try{

        Log(@"--- Testing %3u-bit #%i %s", sizeInBits, (int)algorithm,

            (inKeychain ?", in keychain" :""));

        // Generate key:

        if (inKeychain)

            key = [inKeychain generateSymmetricKeyOfSize: sizeInBits algorithm: algorithm];

        else

            key = [MYSymmetricKey generateSymmetricKeyOfSize: sizeInBits algorithm: algorithm];

        Log(@"Created %@", key);

    CAssert(key);

        CAssertEq(key.algorithm, algorithm);

        CAssertEq(key.keySizeInBits, sizeInBits);

    #if !TARGET_OS_IPHONE

        CAssert(key.cssmKey != NULL);

    #endif

        NSData *keyData = key.keyData;

        Log(@"Key data = %@", keyData);

        CAssertEq(keyData.length, sizeInBits/8);

        // Encrypt a small amount of text:

        Log(@"Testing encryption / decryption ...");

        NSData *cleartext = [@"This is a test. This is only a test." dataUsingEncoding: NSUTF8StringEncoding];

        NSData *encrypted = [key encryptData: cleartext];

        Log(@"Encrypted = %u bytes: %@", encrypted.length, encrypted);

        CAssert(encrypted.length >= cleartext.length);

        NSData *decrypted = [key decryptData: encrypted];

        CAssertEqual(decrypted, cleartext);

        // Encrypt large binary data:

        cleartext = [NSData dataWithContentsOfFile: @"/Library/Desktop Pictures/Nature/Zen Garden.jpg"];

        CAssert(cleartext);

        encrypted = [key encryptData: cleartext];

        Log(@"Encrypted = %u bytes", encrypted.length);

        CAssert(encrypted.length >= cleartext.length);

        decrypted = [key decryptData: encrypted];

        CAssertEqual(decrypted, cleartext);

    #if 1

        Log(@"Testing initWithKeyData:...");

        MYSymmetricKey *key2 = [[MYSymmetricKey alloc] initWithKeyData: keyData algorithm: algorithm];

        CAssert(key2);

        Log(@"Key from data = %@",key2);

        CAssertEqual(key2.keyData, keyData);

        CAssertEq(key2.algorithm, algorithm);

        CAssertEq(key2.keySizeInBits, sizeInBits);

        decrypted = [key2 decryptData: encrypted];

        CAssertEqual(decrypted, cleartext);

        [key2 release];

    #endif

    #if !TARGET_OS_IPHONE

#if 0 // TEMPORARILY OUT OF ORDER

        // Try exporting and importing a wrapped key:

        Log(@"Testing export/import...");

        NSData *exported = [key exportWrappedKeyWithPassphrasePrompt: @"Export symmetric key with passphrase:"];

        Log(@"Exported key: %@", exported);

    #if 1

        CAssert(exported);

    #else

        //FIX: Exporting symmetric keys isn't working. Temporarily making this optional.

        if (!exported)

            Warn(@"Unable to export wrapped key");

        else

    #endif

        {

            CAssert(exported);

            MYSymmetricKey *key2 = [[MYSymmetricKey alloc] initWithKeyData: exported algorithm: algorithm];

            Log(@"Reconstituted as %@", key2);

            CAssertEqual(key2.keyData,key.keyData);

            decrypted = [key2 decryptData: encrypted];

            CAssertEqual(decrypted, cleartext);

        }

#endif 0

    #endif

    }@finally{

        [key removeFromKeychain];

    }

    [pool drain];

}

TestCase(MYSymmetricKey) {

    #define kNTests 11

    static const CCAlgorithm kTestAlgorithms[kNTests] = {

        kCCAlgorithmAES128, kCCAlgorithmAES128, kCCAlgorithmAES128,

        kCCAlgorithmDES, kCCAlgorithm3DES,

        kCCAlgorithmCAST, kCCAlgorithmCAST, kCCAlgorithmCAST,

        kCCAlgorithmRC4, kCCAlgorithmRC4, kCCAlgorithmRC4};

    static const unsigned kTestBitSizes[kNTests] = {

        128, 192, 256,

        64, 3*64,

        40, 80, 128,

        32, 200, 512*8};

    for (int useKeychain=0; useKeychain<=1; useKeychain++)

        for (int testNo=0; testNo<kNTests; testNo++) 

            testSymmetricKey(kTestAlgorithms[testNo], 

                             kTestBitSizes[testNo],

                             useKeychain ?[MYKeychain defaultKeychain] :nil);

}

#if !TARGET_OS_IPHONE

TestCase(MYSymmetricKeyPassphrase) {

    Log(@"Prompting for raw passphrase --");

    NSString *rawPassphrase = [MYSymmetricKey promptForPassphraseWithAlertTitle: @"Raw Passphrase Test" 

                                                                    alertPrompt: @"Enter the passphrase 'Testing':"

                                                                       creating: YES];

    Log(@"You entered: '%@'", rawPassphrase);

    CAssertEqual(rawPassphrase, @"Testing");

    Log(@"Prompting for passphrase for key --");

    MYSymmetricKey *key = [MYSymmetricKey generateFromUserPassphraseWithAlertTitle: @"Symmetric Key Passphrase Test Case" 

                                                                       alertPrompt: @"Please enter a passphrase to generate a key:"

                                                                          creating: YES

                                                                              salt: @"wahooma"];

    Log(@"Key from passphrase = %@", key);

    CAssert(key);

    // Encrypt a small amount of text:

    Log(@"Testing encryption / decryption ...");

    NSData *cleartext = [@"This is a test. This is only a test." dataUsingEncoding: NSUTF8StringEncoding];

    NSData *encrypted = [key encryptData: cleartext];

    Log(@"Encrypted = %u bytes: %@", encrypted.length, encrypted);

    CAssert(encrypted.length >= cleartext.length);

    NSData *decrypted = [key decryptData: encrypted];

    CAssertEqual(decrypted, cleartext);

    // Now test decryption by re-entered passphrase:

    Log(@"Testing decryption using re-entered passphrase...");

    MYSymmetricKey *key2 = [MYSymmetricKey generateFromUserPassphraseWithAlertTitle: @"Symmetric Key Passphrase Test Case" 

                                                                        alertPrompt: @"Please re-enter the same passphrase:" 

                                                                           creating: NO

                                                                               salt: @"wahooma"];

    Log(@"Key from passphrase = %@", key2);

    CAssert(key2);

    decrypted = [key2 decryptData: encrypted];

    CAssertEqual(decrypted, cleartext);

}

#endif

#pragma mark -

#pragma mark KEY-PAIRS:

static void TestUseKeyPair(MYPrivateKey *pair) {

    Log(@"---- TestUseKeyPair { %@ , %@ }.", pair, pair.publicKey);

    CAssert(pair);

    CAssert(pair.keyRef);

    MYPublicKey *publicKey = pair.publicKey;

    CAssert(publicKey.keyRef);

    NSData *pubKeyData = publicKey.keyData;

    Log(@"Public key = %@ (%u bytes)",pubKeyData,pubKeyData.length);

    CAssert(pubKeyData);

    MYSHA1Digest *pubKeyDigest = publicKey.publicKeyDigest;

    Log(@"Public key digest = %@",pubKeyDigest);

    CAssertEqual(pair.publicKeyDigest, pubKeyDigest);

    Log(@"SHA1 of pub key = %@", pubKeyData.my_SHA1Digest.asData);

    CAssertEqual(pubKeyData.my_SHA1Digest,pubKeyDigest);

    // Let's sign data:

    NSData *data = [@"This is a test. This is only a test!" dataUsingEncoding: NSUTF8StringEncoding];

    NSData *sig = [pair signData: data];

    Log(@"Signature = %@ (%u bytes)",sig,sig.length);

    CAssert(sig);

    CAssert( [publicKey verifySignature: sig ofData: data] );

    // Now let's encrypt...

    NSData *crypted = [publicKey rawEncryptData: data];

    Log(@"Encrypted = %@ (%u bytes)",crypted,crypted.length);

    CAssert(crypted);

    CAssertEqual([pair rawDecryptData: crypted], data);

    Log(@"Verified decryption.");

    // Test creating a standalone public key:

    MYPublicKey *pub = [[MYPublicKey alloc] initWithKeyRef: publicKey.keyRef];

    CAssert( [pub verifySignature: sig ofData: data] );

    Log(@"Verified signature.");

    // Test creating a public key from data:

    Log(@"Reconstituting public key from data...");

    pub = [[MYPublicKey alloc] initWithKeyData: pubKeyData];

    CAssert(pub);

    CAssertEqual(pub.keyData, pubKeyData);

    CAssertEqual(pub.publicKeyDigest, pubKeyDigest);

    CAssert( [pub verifySignature: sig ofData: data] );

    Log(@"Verified signature from reconstituted key.");

}

static void TestWrapSessionKey( MYPrivateKey *privateKey ) {

#if !TARGET_OS_IPHONE

    MYSymmetricKey *sessionKey = [MYSymmetricKey generateSymmetricKeyOfSize: 128 algorithm:kCCAlgorithmAES128];

    CAssert(sessionKey);

    NSData *cleartext = [@"This is a test. This is only a test." dataUsingEncoding: NSUTF8StringEncoding];

    NSData *encrypted = [sessionKey encryptData: cleartext];

    Log(@"Wrapping session key %@, %@", sessionKey, sessionKey.keyData);

    NSData *wrapped = [privateKey.publicKey wrapSessionKey: sessionKey];

    Log(@"Wrapped session key = %u bytes: %@", wrapped.length,wrapped);

    CAssert(wrapped.length >= 128/8);

    MYSymmetricKey *unwrappedKey = [privateKey unwrapSessionKey: wrapped

                                                  withAlgorithm: kCCAlgorithmAES128

                                                     sizeInBits: 128];

    Log(@"Unwrapped session key = %@, %@", unwrappedKey, unwrappedKey.keyData);

    CAssert(unwrappedKey);

    CAssertEq(unwrappedKey.algorithm, sessionKey.algorithm);

    CAssertEq(unwrappedKey.keySizeInBits, sessionKey.keySizeInBits);

    CAssertEqual(unwrappedKey.keyData, sessionKey.keyData);

    Log(@"Verifying that unwrapped key works");

    NSData *decrypted = [unwrappedKey decryptData: encrypted];

    CAssertEqual(decrypted, cleartext);

#endif

}

TestCase(MYGenerateKeyPair) {

    RequireTestCase(MYKeychain);

    Log(@"Generating key pair...");

    MYPrivateKey *pair = [[MYKeychain defaultKeychain] generateRSAKeyPairOfSize: kTestCaseRSAKeySize];

    MYPublicKey *publicKey = pair.publicKey;

    Log(@"...created { %@ , %@ }.", pair, publicKey);

    @try{

        TestUseKeyPair(pair);

        TestWrapSessionKey(pair);

        [pair setName: @"Test KeyPair Label"];

        CAssertEqual(pair.name, @"Test KeyPair Label");

        CAssertEqual(publicKey.name, @"Test KeyPair Label");

#if !TARGET_OS_IPHONE

        [pair setComment: @"This key-pair was generated automatically by a test case."];

        CAssertEqual(pair.comment, @"This key-pair was generated automatically by a test case.");

        CAssertEqual(publicKey.comment, @"This key-pair was generated automatically by a test case.");

#endif

        [pair setAlias: @"TestCase@mooseyard.com"];

        CAssertEqual(pair.alias, @"TestCase@mooseyard.com");

        CAssertEqual(publicKey.alias, @"TestCase@mooseyard.com");

        CAssert([pair removeFromKeychain]);

        Log(@"Removed key-pair.");

        pair = nil;

    }@finally {

        if (pair) {

            if ([pair removeFromKeychain])

                Log(@"Removed key-pair from keychain.");

            else

                Warn(@"Unable to remove test key-pair from keychain");

        }

    }

}

#if !TARGET_OS_IPHONE

TestCase(MYUseIdentity) {

    MYIdentity *me = nil;//[MYIdentity preferredIdentityForName: @"MYCryptoTest"];

    if (!me) {

        NSArray *idents = [[[MYKeychain allKeychains] enumerateIdentities] allObjects];

        SFChooseIdentityPanel *panel = [SFChooseIdentityPanel sharedChooseIdentityPanel];

        [panel setAlternateButtonTitle: @"Cancel"];

        if ([panel my_runModalForIdentities: idents 

                                    message: @"Choose an identity for the MYEncoder test case:"]

            != NSOKButton) {

            [NSException raise: NSGenericException format: @"User canceled"];

        }

        me = [panel my_identity];

        [me makePreferredIdentityForName: @"MYCryptoTest"];

    }

    CAssert(me,@"No default identity has been set up in the Keychain");

    TestUseKeyPair(me.privateKey);

}

#endif

#pragma mark -

#pragma mark KEYPAIR EXPORT:

static void testKeyPairExportWithPrompt(BOOL withPrompt) {

    MYKeychain *keychain = [MYKeychain allKeychains];

    Log(@"Generating key pair...");

    MYPrivateKey *pair = [keychain generateRSAKeyPairOfSize: kTestCaseRSAKeySize];

    CAssert(pair);

    CAssert(pair.keyRef);

    CAssert(pair.publicKey.keyRef);

    Log(@"...created pair.");

    @try{

        NSData *pubKeyData = pair.publicKey.keyData;

        CAssert(pubKeyData.length >= kTestCaseRSAKeySize/8);

        [pair setName: @"Test KeyPair Label"];

        CAssertEqual(pair.name, @"Test KeyPair Label");

#if !TARGET_OS_IPHONE

        [pair setComment: @"This key-pair was generated automatically by a test case."];

        CAssertEqual(pair.comment, @"This key-pair was generated automatically by a test case.");

#endif

        [pair setAlias: @"TestCase@mooseyard.com"];

        CAssertEqual(pair.alias, @"TestCase@mooseyard.com");

#if !TARGET_OS_IPHONE

        Log(@"Exporting key-pair...");

        NSString *passphrase = @"passphrase";

        NSData *privKeyData;

        if (withPrompt)

            privKeyData = [pair exportKey];

        else

            privKeyData = [pair _exportKeyInFormat: kSecFormatWrappedOpenSSL

                                          withPEM: YES

                                       passphrase: passphrase];

        Log(@"Exported data = %@ (%u bytes)", privKeyData,privKeyData.length);

        CAssert(privKeyData);

        [privKeyData writeToFile: @"ExportedPrivKey" atomically: YES];

#endif

        // Check key lookup:

        Log(@"Looking up public key of pair in keychain...");

        MYSHA1Digest *digest = pair.publicKeyDigest;

        MYPublicKey *foundKey = [keychain publicKeyWithDigest: digest];

        CAssertEqual(foundKey, pair.publicKey);

        CAssert([keychain.enumeratePublicKeys.allObjects containsObject: pair.publicKey]);

        MYPrivateKey *foundPair = [keychain privateKeyWithDigest: digest];

        CAssertEqual(foundPair, pair);

        CAssert([keychain.enumeratePrivateKeys.allObjects containsObject: pair]);

        Log(@"Removing key-pair from keychain...");

        CAssert([pair removeFromKeychain]);

        pair = nil;

        CAssert([keychain publicKeyWithDigest: digest] == nil);

#if !TARGET_OS_IPHONE

        Log(@"Importing key-pair...");

        if (withPrompt) {

            pair = [keychain importPublicKey: pubKeyData 

                                  privateKey: privKeyData];

        } else {

            pair = [[[MYPrivateKey alloc] _initWithKeyData: privKeyData

                                             publicKeyData: pubKeyData

                                               forKeychain: keychain.keychainRefOrDefault

                                                passphrase: passphrase]

                    autorelease];

        }

        CAssert(pair);

        CAssertEqual(pair.publicKey.keyData, pubKeyData);

#endif

    }@finally {

        if (pair) {

            if ([pair removeFromKeychain])

                Log(@"Removed key-pair from keychain.");

            else

                Warn(@"Unable to remove test key-pair from keychain");

        }

    }

}

TestCase(KeyPairExport) {

    RequireTestCase(MYKeychain);

    RequireTestCase(MYGenerateKeyPair);

    testKeyPairExportWithPrompt(NO);

}

TestCase(KeyPairExportWithUI) {

    RequireTestCase(KeyPairExport);

    testKeyPairExportWithPrompt(YES);

}

#endif DEBUG