//

 //  Cryptor.m

 //  MYCrypto

 //

 //  Created by Jens Alfke on 3/21/09.

 //  Copyright 2009 Jens Alfke. All rights reserved.

 //

 #import "MYCryptor.h"

 #import "MYDigest.h"

 #import "Test.h"

 #if MYCRYPTO_USE_IPHONE_API

 #import <Security/SecRandom.h>

 #else

 #import "MYCrypto_Private.h"

 #import "MYKeychain.h"

 #import <stdlib.h>

 #endif

 NSString* const CryptorErrorDomain = @"CCCryptor";

 #if !MYCRYPTO_USE_IPHONE_API

 static BOOL generateRandomBytes(CSSM_CSP_HANDLE module, uint32_t lengthInBytes, void *dstBytes);

 #endif

 @interface MYCryptor ()

 @property (readwrite, retain) NSError *error;

 @end

 @implementation MYCryptor

 + (NSData*) randomKeyOfLength: (size_t)lengthInBits {

     size_t lengthInBytes = (lengthInBits + 7)/8;

     NSParameterAssert(lengthInBytes<100000);

     uint8_t *bytes = malloc(lengthInBytes);

     if (!bytes) return nil;

 #if MYCRYPTO_USE_IPHONE_API

     BOOL ok = SecRandomCopyBytes(kSecRandomDefault, lengthInBytes,bytes) >= 0;

 #else

     BOOL ok = generateRandomBytes([[MYKeychain defaultKeychain] CSPHandle], lengthInBytes, bytes);

 #endif

     if (ok)

         return [NSData dataWithBytesNoCopy: bytes length: lengthInBytes freeWhenDone: YES];

     else {

         free(bytes);

         return nil;

     }

 }

 + (NSData*) keyOfLength: (size_t)lengthInBits

          fromPassphrase: (NSString*)passphrase

                    salt: (id)salt

 {

     // This follows algorithm PBKDF1 from PKCS#5 v2.0, with Hash=SHA-256 and c=13.

     Assert(passphrase);

     Assert(salt);

     passphrase = $sprintf(@"MYCrypto|%@|%@", passphrase, salt);

     size_t lengthInBytes = (lengthInBits + 7)/8;

     MYDigest *digest = [[passphrase dataUsingEncoding: NSUTF8StringEncoding] my_SHA256Digest];

     for (int i=0; i<12; i++)

         digest = digest.asData.my_SHA256Digest;

     if (lengthInBytes <= digest.length)

         return [digest.asData subdataWithRange: NSMakeRange(0,lengthInBytes)];

     else

         return nil;

 }

 - (id) initWithKey: (NSData*)key

          algorithm: (CCAlgorithm)algorithm

          operation: (CCOperation)op {

     self = [super init];

     if (self) {

         NSParameterAssert(key);

         _key = [key copy];

         _operation = op;

         _algorithm = algorithm;

         _options = kCCOptionPKCS7Padding;

     }

     return self;

 }

 - (id) initEncryptorWithKey: (NSData*)key algorithm: (CCAlgorithm)algorithm {

     return [self initWithKey: key algorithm: algorithm operation: kCCEncrypt];

 }

 - (id) initDecryptorWithKey: (NSData*)key algorithm: (CCAlgorithm)algorithm {

     return [self initWithKey: key algorithm: algorithm operation: kCCDecrypt];

 }

 - (void) dealloc

 {

     if (_cryptor)

         CCCryptorRelease(_cryptor);

     [_key autorelease];

     [_output autorelease];

     [_outputStream release];

     [super dealloc];

 }

 - (void) finalize

 {

     if (_cryptor)

         CCCryptorRelease(_cryptor);

     [super finalize];

 }

 @synthesize key=_key, algorithm=_algorithm, options=_options,

     outputStream=_outputStream, error=_error;

 - (BOOL) _check: (CCCryptorStatus)status {

     if (status == kCCSuccess)

         return YES;

     else {

         Warn(@"MYCryptor: CCCryptor error %i", status);

         self.error = [NSError errorWithDomain: CryptorErrorDomain code: status userInfo: nil];

         return NO;

     }

 }

 - (BOOL) _outputBytes: (const void*)bytes length: (size_t)length {

     if (_outputStream) {

         size_t written = [_outputStream write: bytes maxLength: length];

         if (written < 0) {

             self.error = _outputStream.streamError;

             if (_error)

                 Warn(@"MYCryptor: NSOutputStream error %@", _error);

             else

                 [self _check: kMYCryptorErrorOutputStreamChoked];

             return NO;

         } else if (written < length) {

             [self _check: kMYCryptorErrorOutputStreamChoked];

             return NO;

         }

     } else if (length > 0) {

         [_output appendBytes: bytes length: length];

     }

     return YES;

 }

 - (BOOL) _start {

     if (!_cryptor && !_error) {

         if ([self _check: CCCryptorCreate(_operation, _algorithm, _options,

                                           _key.bytes, _key.length, NULL, &_cryptor)]) {

             _output = [[NSMutableData alloc] initWithCapacity: 1024];

         }

     }

     return !_error;

 }

 - (BOOL) addBytes: (const void*)bytes length: (size_t)length {

     if (length > 0) {

         NSParameterAssert(bytes!=NULL);

         if(!_error && (_cryptor || [self _start])) {

             size_t outputLength = CCCryptorGetOutputLength(_cryptor,length,false);

             void *output = malloc(outputLength);

             if ([self _check: CCCryptorUpdate(_cryptor, bytes, length,

                                               output, outputLength, &outputLength)]) {

                 [self _outputBytes: output length: outputLength];

             }

             free(output);

         }

     }

     return !_error;

 }

 - (BOOL) addData: (NSData*)data

 {

     return [self addBytes: data.bytes length: data.length];

 }

 - (BOOL) addString: (NSString*)str {

     return [self addData: [str dataUsingEncoding: NSUTF8StringEncoding]];

 }

 - (BOOL) addFromStream: (NSInputStream*)input

 {

     uint8_t inputBuffer[1024];

     size_t avail;

     while (!_error && input.hasBytesAvailable) {

         avail = sizeof(inputBuffer);

         NSInteger nRead = [input read: inputBuffer maxLength: sizeof(inputBuffer)];

         if (nRead < 0) {

             self.error = input.streamError;

             return NO;

         } else if (nRead == 0) {

             break;

         } else if (![self addBytes: inputBuffer length: nRead])

             return NO;

     }

     return YES;

 }

 - (BOOL) finish

 {

     if(!_error && (_cryptor || [self _start])) {

         size_t outputLength = 100; //CCCryptorGetOutputLength(_cryptor,1,true);

         void *output = malloc(outputLength);

         if ([self _check: CCCryptorFinal(_cryptor, output, outputLength, &outputLength)]) {

             [self _outputBytes: output length: outputLength];

         }

         free(output);

     }

     CCCryptorRelease(_cryptor);

     _cryptor = NULL;

     return !_error;

 }

 - (NSData*) outputData {

     if (_cryptor) [self finish];

     if(_error) {

         [_output release];

         _output = nil;

     }

     return _output;

 }

 - (NSString*) outputString {

     NSData *output = self.outputData;

     if (output) {

         NSString *str = [[NSString alloc] initWithData: output encoding: NSUTF8StringEncoding];

         return [str autorelease];

     } else

         return nil;

 }

 // NSStream delegate method

 - (void)stream:(NSStream *)stream handleEvent:(NSStreamEvent)eventCode {

     switch (eventCode) {

         case NSStreamEventHasBytesAvailable:

             [self addFromStream: (NSInputStream*)stream];

             break;

         case NSStreamEventEndEncountered:

             [self finish];

             break;

         case NSStreamEventErrorOccurred:

             if (!_error)

                 self.error = stream.streamError;

             break;

         default:

             break;

     }

 }

 @end

 #if !MYCRYPTO_USE_IPHONE_API

 static BOOL generateRandomBytes(CSSM_CSP_HANDLE module, uint32_t lengthInBytes, void *dstBytes) {

     // Adapted from code in Keychain.framework's KeychainUtils.m by Wade Tregaskis.

     CSSM_CC_HANDLE ccHandle;

     if (!checkcssm(CSSM_CSP_CreateRandomGenContext(module, CSSM_ALGID_APPLE_YARROW, NULL,

                                                   lengthInBytes, &ccHandle),

                    @"CSSM_CSP_CreateRandomGenContext"))

         return NO;

     CSSM_DATA data = {.Data=dstBytes, .Length=lengthInBytes};

     BOOL ok = checkcssm(CSSM_GenerateRandom(ccHandle, &data), @"CSSM_GenerateRandom");

     CSSM_DeleteContext(ccHandle);

     return ok;

 }

 #endif

 TestCase(MYCryptor) {

     // Encryption:

     NSData *key = [MYCryptor randomKeyOfLength: 256];

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

     MYCryptor *enc = [[MYCryptor alloc] initEncryptorWithKey: key algorithm: kCCAlgorithmAES128];

     CAssert(enc);

     CAssert([enc addString: @"This is a test. "]);

     CAssert([enc addString: @"This is only a test."]);

     CAssertEqual(enc.error, nil);

     NSData *encrypted = enc.outputData;

     CAssertEqual(enc.error, nil);

     CAssert(encrypted.length > 0);

     [enc release];

     Log(@"Encrypted = %@", encrypted);

     // Decryption:

     MYCryptor *dec = [[MYCryptor alloc] initDecryptorWithKey: key algorithm: kCCAlgorithmAES128];

     CAssert(dec);

     CAssert([dec addData: encrypted]);

     NSString *decrypted = dec.outputString;

     CAssertEqual(dec.error, nil);

     [dec release];

     Log(@"Decrypted = '%@'", decrypted);

     CAssertEqual(decrypted, @"This is a test. This is only a test.");

     // Encryption to stream:

     key = [MYCryptor randomKeyOfLength: 256];

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

     enc = [[MYCryptor alloc] initEncryptorWithKey: key algorithm: kCCAlgorithmAES128];

     CAssert(enc);

     enc.outputStream = [NSOutputStream outputStreamToMemory];

     [enc.outputStream open];

     CAssert([enc addString: @"This is a test. "]);

     CAssert([enc addString: @"This is only a test."]);

     CAssert([enc finish]);

     CAssertEqual(enc.error, nil);

     encrypted = [[enc.outputStream propertyForKey: NSStreamDataWrittenToMemoryStreamKey] retain];

     CAssert(encrypted.length > 0);

     [enc release];

     Log(@"Encrypted = %@", encrypted);

     dec = [[MYCryptor alloc] initDecryptorWithKey: key algorithm: kCCAlgorithmAES128];

     CAssert(dec);

     dec.outputStream = [NSOutputStream outputStreamToMemory];

     [dec.outputStream open];

     CAssert([dec addData: encrypted]);

     CAssert([dec finish]);

     CAssertEqual(dec.error, nil);

     NSData *decryptedData = [dec.outputStream propertyForKey: NSStreamDataWrittenToMemoryStreamKey];

     [dec release];

     decrypted = [[NSString alloc] initWithData: decryptedData

                                                       encoding: NSUTF8StringEncoding];

     Log(@"Decrypted = '%@'", decrypted);

     CAssertEqual(decrypted, @"This is a test. This is only a test.");

     [encrypted release];

     [decrypted release];

 }

 /*

  Copyright (c) 2009, Jens Alfke <jens@mooseyard.com>. All rights reserved.

  Redistribution and use in source and binary forms, with or without modification, are permitted

  provided that the following conditions are met:

  * Redistributions of source code must retain the above copyright notice, this list of conditions

  and the following disclaimer.

  * Redistributions in binary form must reproduce the above copyright notice, this list of conditions

  and the following disclaimer in the documentation and/or other materials provided with the

  distribution.

  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR

  IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND 

  FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRI-

  BUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES

  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 

   PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 

  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF 

  THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

  */