//

 //  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"

 #import "MYCrypto+Cocoa.h"

 #import "MYCrypto_Private.h"

 #if DEBUG

 #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 ) {

     NSAutoreleasePool *pool = [NSAutoreleasePool new];

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

     // Generate key:

     MYSymmetricKey *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:

     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 !TARGET_OS_IPHONE

     // Try reconstituting the key from its data:

     NSData *exported = [key exportKeyInFormat: kSecFormatWrappedPKCS8 withPEM: NO];

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

     // CAssert(exported);

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

     if (exported) {

         CAssert(exported);

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

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

         CAssertEqual(key2,key);

         decrypted = [key2 decryptData: encrypted];

         CAssertEqual(decrypted, cleartext);

     } else

         Warn(@"Unable to export key in PKCS8");

 #endif

     [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 i=0; i<kNTests; i++) 

         testSymmetricKey(kTestAlgorithms[i], kTestBitSizes[i]);

 }

 #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);

     // 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 encryptData: data];

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

     CAssert(crypted);

     CAssertEqual([pair decryptData: 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.");

 }

 TestCase(MYGenerateKeyPair) {

     RequireTestCase(MYKeychain);

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

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

     MYPublicKey *publicKey = pair.publicKey;

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

     @try{

         TestUseKeyPair(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");

         }

     }

 }

 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);

 }

 #pragma mark -

 #pragma mark KEYPAIR EXPORT:

 static void testKeyPairExportWithPrompt(BOOL withPrompt) {

     MYKeychain *keychain = [MYKeychain allKeychains];

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

     MYPrivateKey *pair = [keychain generateRSAKeyPairOfSize: 512];

     CAssert(pair);

     CAssert(pair.keyRef);

     CAssert(pair.publicKey.keyRef);

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

     @try{

         NSData *pubKeyData = pair.publicKey.keyData;

         CAssert(pubKeyData.length >= 512/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(MYPrivateKey);

     testKeyPairExportWithPrompt(NO);

 }

 TestCase(KeyPairExportWithUI) {

     RequireTestCase(KeyPairExport);

     testKeyPairExportWithPrompt(YES);

 }

 #endif DEBUG