//

//  MYDEREncoder.m

//  MYCrypto

//

//  Created by Jens Alfke on 5/29/09.

//  Copyright 2009 Jens Alfke. All rights reserved.

//

// Reference:

// <http://www.columbia.edu/~ariel/ssleay/layman.html> "Layman's Guide To ASN.1/BER/DER"

#import "MYDEREncoder.h"

#import "MYASN1Object.h"

#import "MYBERParser.h"

#import "MYOID.h"

#import "MYErrorUtils.h"

#define MYDEREncoderException @"MYDEREncoderException"

@interface MYDEREncoder ()

- (void) _encode: (id)object;

@property (retain) NSError *error;

/* Forces use of PrintableString tag for ASCII strings that contain characters not valid

    for that encoding (notably '@'). Provided to get byte-for-byte compatibility with certs

    generated by CDSA, for test cases that check this. */

@property BOOL _forcePrintableStrings;

@end

@implementation MYDEREncoder

- (id) initWithRootObject: (id)rootObject

{

    self = [super init];

    if (self != nil) {

        _rootObject = [rootObject retain];

    }

    return self;

}

+ (NSData*) encodeRootObject: (id)rootObject error: (NSError**)outError {

    MYDEREncoder *encoder = [[self alloc] initWithRootObject: rootObject];

    NSData *output = [encoder.output copy];

    if (outError) *outError = [[encoder.error retain] autorelease];

    [encoder release];

    return [output autorelease];

}

- (void) dealloc

{

    [_rootObject release];

    [_output release];

    [_error release];

    [super dealloc];

}

- (id) copyWithZone: (NSZone*)zone {

    MYDEREncoder *copy = [[[self class] alloc] init];

    copy->_forcePrintableStrings = _forcePrintableStrings;

    return copy;

}

static unsigned sizeOfUnsignedInt (UInt64 n) {

    unsigned bytes = 0;

    for (; n; n >>= 8)

        bytes++;

    return bytes;

}

static unsigned encodeUnsignedInt (UInt64 n, UInt8 buf[], BOOL padHighBit) {

    unsigned size = MAX(1U, sizeOfUnsignedInt(n));

    UInt64 bigEndian = NSSwapHostLongLongToBig(n);

    const UInt8* src = (UInt8*)&bigEndian + (8-size);

    UInt8 *dst = &buf[0];

    if (padHighBit && (*src & 0x80)) {

        *dst++ = 0;

        size++;

    }

    memcpy(dst, src, size);

    return size;

}

static unsigned encodeSignedInt (SInt64 n, UInt8 buf[]) {

    if (n >= 0)

        return encodeUnsignedInt(n, buf, YES);

    else {

        unsigned size = MAX(1U, sizeOfUnsignedInt(~n));

        UInt64 bigEndian = NSSwapHostLongLongToBig(n);

        const UInt8* src = (UInt8*)&bigEndian + (8-size);

        UInt8 *dst = &buf[0];

        if (!(*src & 0x80)) {

            *dst++ = 0xFF;

            size++;

        }

        memcpy(dst, src, size);

        return size;

    }

}

- (void) _writeTag: (unsigned)tag

             class: (unsigned)tagClass

       constructed: (BOOL) constructed

            length: (size_t)length 

{

    struct {

        unsigned tag            :5;

        unsigned isConstructed  :1;

        unsigned tagClass       :2;

        unsigned length         :7;

        unsigned isLengthLong   :1;

        UInt8    extraLength[9];

    } header;

    size_t headerSize = 2;

    header.tag = tag;

    header.isConstructed = constructed;

    header.tagClass = tagClass;

    if (length < 128) {

        header.isLengthLong = NO;

        header.length = length;

    } else {

        header.isLengthLong = YES;

        header.length = encodeUnsignedInt(length, header.extraLength, NO);

        headerSize += header.length;

    }

    [_output appendBytes: &header length: headerSize];

}

- (void) _writeTag: (unsigned)tag

             class: (unsigned)tagClass

       constructed: (BOOL) constructed

             bytes: (const void*)bytes 

            length: (size_t)length 

{

    [self _writeTag: tag class: tagClass constructed: constructed length: length];

    [_output appendBytes: bytes length: length];

}

- (void) _writeTag: (unsigned)tag

             class: (unsigned)tagClass

       constructed: (BOOL) constructed

              data: (NSData*)data 

{

    Assert(data);

    [self _writeTag: tag class: tagClass constructed: constructed bytes: data.bytes length: data.length];

}

- (void) _encodeNumber: (NSNumber*)number {

    // Special-case detection of booleans by pointer equality, because otherwise they appear

    // identical to 0 and 1:

    if (number==$true || number==$false) {

        UInt8 value = number==$true ?0xFF :0x00;

        [self _writeTag: 1 class: 0 constructed: NO bytes: &value length: 1];

        return;

    }

    const char *type = number.objCType;

    if (strlen(type) == 1) {

        switch(type[0]) {

            case 'c':

            case 'i':

            case 's':

            case 'l':

            case 'q':

            {   // Signed integers:

                UInt8 buf[9];

                size_t size = encodeSignedInt(number.longLongValue, buf);

                [self _writeTag: 2 class: 0 constructed: NO bytes: buf length: size];

                return;

            }

            case 'C':

            case 'I':

            case 'S':

            case 'L':

            case 'Q':

            {   // Unsigned integers:

                UInt8 buf[9];

                size_t size = encodeUnsignedInt(number.unsignedLongLongValue, buf, YES);

                [self _writeTag: 2 class: 0 constructed: NO bytes: buf length: size];

                return;

            }

            case 'B':

            {   // bool

                UInt8 value = number.boolValue ?0xFF :0x00;

                [self _writeTag: 1 class: 0 constructed: NO bytes: &value length: 1];

                return;

            }

        }

    }

    [NSException raise: MYDEREncoderException format: @"Can't DER-encode value %@ (typecode=%s)", number,type];

}

- (void) _encodeString: (NSString*)string {

    static NSMutableCharacterSet *kNotPrintableCharSet;

    if (!kNotPrintableCharSet) {

        kNotPrintableCharSet = [[NSMutableCharacterSet characterSetWithCharactersInString: @" '()+,-./:=?"] retain];

        [kNotPrintableCharSet formUnionWithCharacterSet: [NSCharacterSet alphanumericCharacterSet]];

        [kNotPrintableCharSet invert];

    }

    NSData *data = [string dataUsingEncoding: NSASCIIStringEncoding];

    if (data) {

        unsigned tag = 19; // printablestring (a silly arbitrary subset of ASCII defined by ASN.1)

        if (!_forcePrintableStrings && [string rangeOfCharacterFromSet: kNotPrintableCharSet].length > 0)

            tag = 20; // IA5string (full 7-bit ASCII)

        [self _writeTag: tag class: 0 constructed: NO data: data];

    } else {

        // fall back to UTF-8:

        [self _writeTag: 12 class: 0 constructed: NO data: [string dataUsingEncoding: NSUTF8StringEncoding]];

    }

}

- (void) _encodeBitString: (MYBitString*)bitString {

    NSUInteger bitCount = bitString.bitCount;

    [self _writeTag: 3 class: 0 constructed: NO length: 1 + (bitCount/8)];

    UInt8 unused = (8 - (bitCount % 8)) % 8;

    [_output appendBytes: &unused length: 1];

    [_output appendBytes: bitString.bits.bytes length: bitCount/8];

}

- (void) _encodeDate: (NSDate*)date {

    NSString *dateStr = [MYBERGeneralizedTimeFormatter() stringFromDate: date];

    [self _writeTag: 24 class: 0 constructed: NO data: [dateStr dataUsingEncoding: NSASCIIStringEncoding]];

}

- (void) _encodeCollection: (id)collection tag: (unsigned)tag class: (unsigned)tagClass {

    MYDEREncoder *subEncoder = [self copy];

    for (id object in collection)

        [subEncoder _encode: object];

    [self _writeTag: tag class: tagClass constructed: YES data: subEncoder.output];

    [subEncoder release];

}

- (void) _encode: (id)object {

    if (!_output)

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

    if ([object isKindOfClass: [NSNumber class]]) {

        [self _encodeNumber: object];

    } else if ([object isKindOfClass: [NSData class]]) {

        [self _writeTag: 4 class: 0 constructed: NO data: object];

    } else if ([object isKindOfClass: [MYBitString class]]) {

        [self _encodeBitString: object];

    } else if ([object isKindOfClass: [NSString class]]) {

        [self _encodeString: object];

    } else if ([object isKindOfClass: [NSDate class]]) {

        [self _encodeDate: object];

    } else if ([object isKindOfClass: [NSNull class]]) {

        [self _writeTag: 5 class: 0 constructed: NO bytes: NULL length: 0];

    } else if ([object isKindOfClass: [NSArray class]]) {

        [self _encodeCollection: object tag: 16 class: 0];

    } else if ([object isKindOfClass: [NSSet class]]) {

        [self _encodeCollection: object tag: 17 class: 0];

    } else if ([object isKindOfClass: [MYOID class]]) {

        [self _writeTag: 6 class: 0 constructed: NO data: [object DEREncoding]];

    } else if ([object isKindOfClass: [MYASN1Object class]]) {

        MYASN1Object *asn = object;

        if (asn.components)

            [self _encodeCollection: asn.components tag: asn.tag class: asn.tagClass];

        else

            [self _writeTag: asn.tag 

                      class: asn.tagClass

                constructed: asn.constructed

                       data: asn.value];

    } else {

        [NSException raise: MYDEREncoderException format: @"Can't DER-encode a %@", [object class]];

    }

}

- (NSData*) output {

    if (!_output && !_error) {

        @try{

            [self _encode: _rootObject];

        }@catch (NSException *x) {

            if ($equal(x.name, MYDEREncoderException)) {

                self.error = MYError(2,MYASN1ErrorDomain, @"%@", x.reason);

                return nil;

            } else

                @throw(x);

        }

    }

    return _output;

}

@synthesize error=_error, _forcePrintableStrings;

@end

#define $data(BYTES...)    ({const uint8_t bytes[] = {BYTES}; [NSData dataWithBytes: bytes length: sizeof(bytes)];})

TestCase(DEREncoder) {

    CAssertEqual([MYDEREncoder encodeRootObject: [NSNull null] error: nil],

                 $data(0x05, 0x00));

    CAssertEqual([MYDEREncoder encodeRootObject: $true error: nil],

                 $data(0x01, 0x01, 0xFF));

    CAssertEqual([MYDEREncoder encodeRootObject: $false error: nil],

                 $data(0x01, 0x01, 0x00));

    // Integers:

    CAssertEqual([MYDEREncoder encodeRootObject: $object(0) error: nil],

                 $data(0x02, 0x01, 0x00));

    CAssertEqual([MYDEREncoder encodeRootObject: $object(1) error: nil],

                 $data(0x02, 0x01, 0x01));

    CAssertEqual([MYDEREncoder encodeRootObject: $object(-1) error: nil],

                 $data(0x02, 0x01, 0xFF));

    CAssertEqual([MYDEREncoder encodeRootObject: $object(72) error: nil],

                  $data(0x02, 0x01, 0x48));

    CAssertEqual([MYDEREncoder encodeRootObject: $object(-128) error: nil],

                 $data(0x02, 0x01, 0x80));

    CAssertEqual([MYDEREncoder encodeRootObject: $object(128) error: nil],

                 $data(0x02, 0x02, 0x00, 0x80));

    CAssertEqual([MYDEREncoder encodeRootObject: $object(255) error: nil],

                 $data(0x02, 0x02, 0x00, 0xFF));

    CAssertEqual([MYDEREncoder encodeRootObject: $object(-256) error: nil],

                 $data(0x02, 0x02, 0xFF, 0x00));

    CAssertEqual([MYDEREncoder encodeRootObject: $object(12345) error: nil],

                 $data(0x02, 0x02, 0x30,0x39));

    CAssertEqual([MYDEREncoder encodeRootObject: $object(-12345) error: nil],

                 $data(0x02, 0x02, 0xCF, 0xC7));

    CAssertEqual([MYDEREncoder encodeRootObject: $object(123456789) error: nil],

                 $data(0x02, 0x04, 0x07, 0x5B, 0xCD, 0x15));

    CAssertEqual([MYDEREncoder encodeRootObject: $object(-123456789) error: nil],

                 $data(0x02, 0x04, 0xF8, 0xA4, 0x32, 0xEB));

    CAssertEqual([MYDEREncoder encodeRootObject: $object(-123456789) error: nil],

                 $data(0x02, 0x04, 0xF8, 0xA4, 0x32, 0xEB));

    // Strings:

    CAssertEqual([MYDEREncoder encodeRootObject: @"hello" error: nil],

                 $data(0x13, 0x05, 'h', 'e', 'l', 'l', 'o'));

    CAssertEqual([MYDEREncoder encodeRootObject: @"thérè" error: nil],

                 $data(0x0C, 0x07, 't', 'h', 0xC3, 0xA9, 'r', 0xC3, 0xA8));

    // Dates:

    CAssertEqual([MYDEREncoder encodeRootObject: [NSDate dateWithTimeIntervalSinceReferenceDate: 265336576]

                                          error: nil],

                 $data(0x18, 0x0F, '2', '0', '0', '9', '0', '5', '3', '0', '0', '0', '3', '6', '1', '6', 'Z'));

    // Sequences:

    CAssertEqual([MYDEREncoder encodeRootObject: $array($object(72), $true) error: nil],

                 $data(0x30, 0x06,  0x02, 0x01, 0x48,  0x01, 0x01, 0xFF));

    CAssertEqual([MYDEREncoder encodeRootObject: $array( $array($object(72), $true), 

                                                         $array($object(72), $true))

                                          error: nil],

                 $data(0x30, 0x10,  

                       0x30, 0x06,  0x02, 0x01, 0x48,  0x01, 0x01, 0xFF,

                       0x30, 0x06,  0x02, 0x01, 0x48,  0x01, 0x01, 0xFF));

}

TestCase(EncodeCert) {

    NSError *error = nil;

    NSData *cert = [NSData dataWithContentsOfFile: @"../../Tests/selfsigned.cer"];

    id certObjects = MYBERParse(cert,&error);

    CAssertNil(error);

    Log(@"Decoded as:\n%@", [MYASN1Object dump: certObjects]);

    MYDEREncoder *encoder = [[MYDEREncoder alloc] initWithRootObject: certObjects];

    encoder._forcePrintableStrings = YES;       // hack for compatibility with the way CDSA writes ASN.1

    NSData *encoded = encoder.output;

    CAssertNil(error);

    id reDecoded = MYBERParse(encoded, &error);

    CAssertNil(error);

    Log(@"Re-decoded as:\n%@", [MYASN1Object dump: reDecoded]);

    [encoded writeToFile: @"../../Tests/selfsigned_reencoded.cer" atomically: YES];

    CAssertEqual(encoded,cert);

}

/*

 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.

 */