//
// Copyright 2011-2012 Kosher Penguin LLC
// Created by Adar Porat (https://github.com/aporat) on 1/16/2012.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#import "KKKeychain.h"
#import <Security/Security.h>
#import "BundleUtil.h"

@implementation KKKeychain

+ (NSString*)appName
{
 NSBundle *bundle = [BundleUtil mainBundle];
 NSString *appName = [bundle objectForInfoDictionaryKey:@"CFBundleDisplayName"];
 if (!appName) {
  appName = [bundle objectForInfoDictionaryKey:@"CFBundleName"];
 }
 return appName;
}

+ (BOOL)setString:(NSString*)string forKey:(NSString*)key
{
 if (string == nil || key == nil) {
  return NO;
 }
 
 key = [NSString stringWithFormat:@"%@ - %@", [KKKeychain appName], key];
    
 // First check if it already exists, by creating a search dictionary and requesting that
 // nothing be returned, and performing the search anyway.
 NSMutableDictionary *existsQueryDictionary = [NSMutableDictionary dictionary];
 
 NSData *data = [string dataUsingEncoding:NSUTF8StringEncoding];
 
 [existsQueryDictionary setObject:(__bridge id)kSecClassGenericPassword forKey:(__bridge id)kSecClass];
 
 // Add the keys to the search dict
 [existsQueryDictionary setObject:@"service" forKey:(__bridge id)kSecAttrService];
 [existsQueryDictionary setObject:key forKey:(__bridge id)kSecAttrAccount];
    
    // Have SecItemCopyMatching return the data even though we're not interested in it (to work around a Citrix Worx bug)
    CFTypeRef dataDummy = nil;
    [existsQueryDictionary setObject:(id)kCFBooleanTrue forKey:(__bridge id)kSecReturnData];
    
    OSStatus res = SecItemCopyMatching((__bridge CFDictionaryRef)existsQueryDictionary, &dataDummy);
    
    if (dataDummy)
        CFRelease(dataDummy);
    
    [existsQueryDictionary removeObjectForKey:(__bridge id)kSecReturnData];
    
 if (res == errSecItemNotFound) {
  if (string != nil) {
   NSMutableDictionary *addDict = existsQueryDictionary;
   [addDict setObject:data forKey:(__bridge id)kSecValueData];
            [addDict setObject:(__bridge id)kSecAttrAccessibleAfterFirstUnlock forKey:(__bridge id)kSecAttrAccessible];
            
   res = SecItemAdd((__bridge CFDictionaryRef)addDict, NULL);
   NSAssert1(res == errSecSuccess, @"Recieved %ld from SecItemAdd!", (long)res);
  }
 } else if (res == errSecSuccess) {
  // Modify an existing one
  // Actually pull it now of the keychain at this point.
        NSDictionary *attributeDict = @{
                                        (__bridge id)kSecValueData: data,
                                        (__bridge id)kSecAttrAccessible: (__bridge id)kSecAttrAccessibleAfterFirstUnlock
                                        };
  res = SecItemUpdate((__bridge CFDictionaryRef)existsQueryDictionary, (__bridge CFDictionaryRef)attributeDict);
  NSAssert1(res == errSecSuccess, @"SecItemUpdated returned %ld!", (long)res);
 } else {
  NSAssert1(NO, @"Received %ld from SecItemCopyMatching!", (long)res);
 }
 return YES;
}

+ (NSString*)getStringForKey:(NSString*)key
{
 key = [NSString stringWithFormat:@"%@ - %@", [KKKeychain appName], key];
 NSMutableDictionary *existsQueryDictionary = [NSMutableDictionary dictionary];
 [existsQueryDictionary setObject:(__bridge id)kSecClassGenericPassword forKey:(__bridge id)kSecClass];
 
 // Add the keys to the search dict
 [existsQueryDictionary setObject:@"service" forKey:(__bridge id)kSecAttrService];
 [existsQueryDictionary setObject:key forKey:(__bridge id)kSecAttrAccount];
 
 // We want the data back!
 CFTypeRef data = nil;
 
 [existsQueryDictionary setObject:(id)kCFBooleanTrue forKey:(__bridge id)kSecReturnData];
    
 OSStatus res = SecItemCopyMatching((__bridge CFDictionaryRef)existsQueryDictionary, &pan style='color:red'>data);
    
 if (res == errSecSuccess) {
  NSString *string = [[NSString alloc] initWithData:(__bridge NSData*)data encoding:NSUTF8StringEncoding];
        CFRelease(data);
  return string;
    } else {
        if (data)
            CFRelease(data);
  NSAssert1(res == errSecItemNotFound, @"SecItemCopyMatching returned %ld!", (long)res);
 }
 
 return nil;
}

+ (void)upgradeAccessibilityForKey:(NSString*)key {
    // Check if the accessibility attributes need to be upgraded
    NSString *derivedKey = [NSString stringWithFormat:@"%@ - %@", [KKKeychain appName], key];
    NSMutableDictionary *queryDictionary = [NSMutableDictionary dictionary];
    [queryDictionary setObject:(__bridge id)kSecClassGenericPassword forKey:(__bridge id)kSecClass];
    
    // Add the keys to the search dict
    [queryDictionary setObject:@"service" forKey:(__bridge id)kSecAttrService];
    [queryDictionary setObject:derivedKey forKey:(__bridge id)kSecAttrAccount];
    
    CFTypeRef attrs = nil;
    
    [queryDictionary setObject:(id)kCFBooleanTrue forKey:(__bridge id)kSecReturnAttributes];
    [queryDictionary setObject:(id)kCFBooleanTrue forKey:(__bridge id)kSecReturnData];
    
    OSStatus res = SecItemCopyMatching((__bridge CFDictionaryRef)queryDictionary, &attrs);
    
    if (res == errSecSuccess) {
        NSDictionary *attrsDict = (__bridge NSDictionary*)attrs;
        if (![[attrsDict objectForKey:(__bridge id)kSecAttrAccessible] isEqualToString:(__bridge NSString*)kSecAttrAccessibleAfterFirstUnlock]) {
            // Need to upgrade the accessibility on this key by setting it again
            [KKKeychain setString:[[NSString alloc] initWithData:[attrsDict objectForKey:(__bridge id)kSecValueData] encoding:NSUTF8StringEncoding] forKey:key];
        }
    }
    
    if (attrs)
        CFRelease(attrs);
}


@end