#import "SNTCodesignChecker.h"

#import <Security/Security.h>

#import "SNTCertificate.h"


// kStaticSigningFlags are the flags used when validating signatures on disk.
//
// Don't validate resources but do validate nested code. Ignoring resources _dramatically_ speeds
// up validation (see below) but does mean images, plists, etc will not be checked and modifying
// these will not be considered invalid. To ensure any code inside the binary is still checked,
// we check nested code.
//
// Timings with different flags:
//   Checking Xcode 5.1.1 bundle:
//      kSecCSDefaultFlags:                                   3.895s
//      kSecCSDoNotValidateResources:                         0.013s
//      kSecCSDoNotValidateResources | kSecCSCheckNestedCode: 0.013s
//
//   Checking Google Chrome 36.0.1985.143 bundle:
//      kSecCSDefaultFlags:                                   0.529s
//      kSecCSDoNotValidateResources:                         0.032s
//      kSecCSDoNotValidateResources | kSecCSCheckNestedCode: 0.033s
//
const SecCSFlags kStaticSigningFlags = kSecCSDoNotValidateResources | kSecCSCheckNestedCode;

// kSigningFlags are the flags used when validating signatures for running binaries.
//
// No special flags needed currently.
const SecCSFlags kSigningFlags = kSecCSDefaultFlags;


@implementation SNTCodesignChecker

#pragma mark Init/dealloc

- (instancetype)initWithSecStaticCodeRef:(SecStaticCodeRef)codeRef {
  self = [super init];

  if (self) {
    _codeRef = codeRef;
    CFRetain(_codeRef);
  }

  return self;
}

- (instancetype)initWithBinaryPath:(NSString *)binaryPath {
  SecStaticCodeRef codeRef = NULL;

  // Get SecStaticCodeRef for binary
  if (SecStaticCodeCreateWithPath((__bridge CFURLRef)[NSURL fileURLWithPath:binaryPath
                                                                isDirectory:NO],
                                  kSecCSDefaultFlags,
                                  &codeRef) == errSecSuccess) {
    self = [self initWithSecStaticCodeRef:codeRef];
  } else {
    self = nil;
  }

  if (codeRef) CFRelease(codeRef);
  return self;
}

- (instancetype)initWithPID:(pid_t)PID {
  SecCodeRef codeRef = NULL;
  NSDictionary *attributes = @{(__bridge NSString *)kSecGuestAttributePid: @(PID)};

  if (SecCodeCopyGuestWithAttributes(NULL,
                                     (__bridge CFDictionaryRef)attributes,
                                     kSecCSDefaultFlags,
                                     &codeRef) == errSecSuccess) {
    self = [self initWithSecStaticCodeRef:codeRef];
  } else {
    self = nil;
  }

  if (codeRef) CFRelease(codeRef);
  return self;
}

- (instancetype)initWithSelf {
  SecCodeRef codeSelf = NULL;
  if (SecCodeCopySelf(kSecCSDefaultFlags, &codeSelf) == errSecSuccess) {
    self = [self initWithSecStaticCodeRef:codeSelf];
  } else {
    self = nil;
  }

  if (codeSelf) CFRelease(codeSelf);
  return self;
}

- (instancetype)init {
  [self doesNotRecognizeSelector:_cmd];
  return nil;
}

- (void)dealloc {
  if (_codeRef) CFRelease(_codeRef);
}

#pragma mark Validate

- (OSStatus)validate {
  return [self validateWithRequirement:NULL];
}

- (OSStatus)validateAppleAnchor {
  SecRequirementRef req = NULL;
  SecRequirementCreateWithString(CFSTR("anchor apple"), kSecCSDefaultFlags, &req);
  return [self validateWithRequirement:req];
}

- (OSStatus)validateAppleAnchorGeneric {
  SecRequirementRef req = NULL;
  SecRequirementCreateWithString(CFSTR("anchor apple generic"), kSecCSDefaultFlags, &req);
  return [self validateWithRequirement:req];
}

- (OSStatus)validateWithRequirement:(SecRequirementRef)requirement {
  // Validate the binary and save the return code.
  if (CFGetTypeID(self.codeRef) == SecStaticCodeGetTypeID()) {
    return SecStaticCodeCheckValidity(self.codeRef, kStaticSigningFlags, requirement);
  } else if (CFGetTypeID(self.codeRef) == SecCodeGetTypeID()) {
    return SecCodeCheckValidity((SecCodeRef)self.codeRef, kSigningFlags, requirement);
  } else {
    return errSecCSSignatureNotVerifiable;
  }
}

#pragma mark Description

- (NSString *)description {
  NSString *retStr;
  if (CFGetTypeID(self.codeRef) == SecStaticCodeGetTypeID()) {
    retStr = @"On-disk binary, ";
  } else {
    retStr = @"In-memory binary, ";
  }

  if ([self validate] == errSecSuccess) {
    [retStr appendFormat:@"signed by %@, ", self.leafCertificate.orgName];
  } else {
    [retStr appendFormat:@"unsigned, "];
  }

  [retStr appendFormat:@"located at: %@", self.binaryPath];

  return retStr;
}

#pragma mark Public accessors

- (NSDictionary *)signingInformation {
  static NSDictionary *signingInformation = nil;
  static dispatch_once_t onceToken;
  dispatch_once(&onceToken, ^{
    // Get dictionary of signing information for binary
    CFDictionaryRef signingDict = NULL;
    SecCodeCopySigningInformation(self.codeRef, kSecCSSigningInformation, &signingDict);
    signingInformation = CFBridgingRelease(signingDict);
  });
  return signingInformation;
}

- (NSArray *)certificates {
  static NSArray *certificates = nil;
  static dispatch_once_t onceToken;
  dispatch_once(&onceToken, ^{
    // Get array of certificates, wrap each one in a SNTCertificate and store in a new array.
    NSArray *certs = self.signingInformation[(__bridge NSString *)kSecCodeInfoCertificates];

    NSMutableArray *tempCerts = [[NSMutableArray alloc] initWithCapacity:certs.count];
    for (id cert in certs) {
      SNTCertificate *newCert =
          [[SNTCertificate alloc] initWithSecCertificateRef:(SecCertificateRef)cert];
      if (newCert) [tempCerts addObject:newCert];
    }
    certificates = [tempCerts copy];
  });
  return certificates;
}

- (SNTCertificate *)leafCertificate {
  return [self.certificates firstObject];
}

- (NSString *)binaryPath {
  CFURLRef path;
  OSStatus status = SecCodeCopyPath(_codeRef, kSecCSDefaultFlags, &path);
  NSURL *pathURL = CFBridgingRelease(path);
  if (status != errSecSuccess) return nil;
  return [pathURL path];
}

#pragma mark Comparisons

- (BOOL)signingChainMatches:(SNTCodesignChecker *)otherChecker {
  return [self.certificates isEqual:otherChecker.certificates];
}

- (BOOL)teamSigningMatches:(SNTCodesignChecker *)otherChecker {
  SNTCertificate *myLeaf = [self.certificates firstObject];
  SNTCertificate *otherLeaf = [otherChecker.certificates firstObject];

  return ([myLeaf.orgUnit isEqual:otherLeaf.orgUnit] &&
          [self validateAppleAnchorGeneric] &&
          [otherChecker validateAppleAnchorGeneric]);
}

@end