#
# OCSP request/response syntax
#
# Derived from a minimal OCSP library (RFC2560) code written by
# Bud P. Bruegger <bud@ancitel.it>
# Copyright: Ancitel, S.p.a,  Rome, Italy
# License: BSD
#

#
# current limitations:
# * request and response works only for a single certificate
# * only some values are parsed out of the response
# * the request does't set a nonce nor signature
# * there is no signature validation of the response
# * dates are left as strings in GeneralizedTime format -- datetime.datetime
# would be nicer
#

import hashlib
from pyasn1.type import tag, namedtype, namedval, univ, constraint,char,useful
from pyasn1.codec.der import decoder, encoder
from pyasn1 import error
import x509

# Start of OCSP module definitions

# This should be in directory Authentication Framework (X.509) module

class CRLReason(univ.Enumerated):
    namedValues = namedval.NamedValues(
        ('unspecified', 0),
        ('keyCompromise', 1),
        ('cACompromise', 2),
        ('affiliationChanged', 3),
        ('superseded', 4),
        ('cessationOfOperation', 5),
        ('certificateHold', 6),
        ('removeFromCRL', 8),
        ('privilegeWithdrawn', 9),
        ('aACompromise', 10)
        )

# end of directory Authentication Framework (X.509) module

# This should be in PKIX Certificate Extensions module

class GeneralName(univ.OctetString): pass

# end of PKIX Certificate Extensions module

id_kp_OCSPSigning = univ.ObjectIdentifier((1, 3, 6, 1, 5, 5, 7, 3, 9))
id_pkix_ocsp = univ.ObjectIdentifier((1, 3, 6, 1, 5, 5, 7, 48, 1))
id_pkix_ocsp_basic = univ.ObjectIdentifier((1, 3, 6, 1, 5, 5, 7, 48, 1, 1))
id_pkix_ocsp_nonce = univ.ObjectIdentifier((1, 3, 6, 1, 5, 5, 7, 48, 1, 2))
id_pkix_ocsp_crl = univ.ObjectIdentifier((1, 3, 6, 1, 5, 5, 7, 48, 1, 3))
id_pkix_ocsp_response = univ.ObjectIdentifier((1, 3, 6, 1, 5, 5, 7, 48, 1, 4))
id_pkix_ocsp_nocheck = univ.ObjectIdentifier((1, 3, 6, 1, 5, 5, 7, 48, 1, 5))
id_pkix_ocsp_archive_cutoff = univ.ObjectIdentifier((1, 3, 6, 1, 5, 5, 7, 48, 1, 6))
id_pkix_ocsp_service_locator = univ.ObjectIdentifier((1, 3, 6, 1, 5, 5, 7, 48, 1, 7))

class AcceptableResponses(univ.SequenceOf):
    componentType = univ.ObjectIdentifier()

class ArchiveCutoff(useful.GeneralizedTime): pass

class UnknownInfo(univ.Null): pass

class RevokedInfo(univ.Sequence):
    componentType = namedtype.NamedTypes(
        namedtype.NamedType('revocationTime', useful.GeneralizedTime()),
        namedtype.OptionalNamedType('revocationReason', CRLReason().subtype(explicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatSimple, 0)))
        )

class CertID(univ.Sequence):
    componentType = namedtype.NamedTypes(
        namedtype.NamedType('hashAlgorithm', x509.AlgorithmIdentifier()),
        namedtype.NamedType('issuerNameHash', univ.OctetString()),
        namedtype.NamedType('issuerKeyHash', univ.OctetString()),
        namedtype.NamedType('serialNumber', x509.CertificateSerialNumber())
        )

class CertStatus(univ.Choice):
    componentType = namedtype.NamedTypes(
        namedtype.NamedType('good', univ.Null().subtype(implicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatSimple, 0))),
        namedtype.NamedType('revoked', RevokedInfo().subtype(implicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatSimple, 1))),
        namedtype.NamedType('unknown', UnknownInfo().subtype(implicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatSimple, 2)))
        )

class SingleResponse(univ.Sequence):
    componentType = namedtype.NamedTypes(
        namedtype.NamedType('certID', CertID()),
        namedtype.NamedType('certStatus', CertStatus()),
        namedtype.NamedType('thisUpdate', useful.GeneralizedTime()),
        namedtype.OptionalNamedType('nextUpdate', useful.GeneralizedTime().subtype(explicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatSimple, 0))),
        namedtype.OptionalNamedType('singleExtensions', x509.Extensions().subtype(explicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatSimple, 1)))
        )

class KeyHash(univ.OctetString): pass

class ResponderID(univ.Choice):
    componentType = namedtype.NamedTypes(
        namedtype.NamedType('byName', x509.Name().subtype(implicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatSimple, 1))),
        namedtype.NamedType('byKey', KeyHash().subtype(implicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatSimple, 2)))
        )

class Version(univ.Integer):
    namedValues = namedval.NamedValues(('v1', 0))

class ResponseData(univ.Sequence):
    componentType = namedtype.NamedTypes(
        namedtype.DefaultedNamedType('version', Version('v1').subtype(explicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatSimple, 0))),
        namedtype.NamedType('responderID', ResponderID()),
        namedtype.NamedType('producedAt', useful.GeneralizedTime()),
        namedtype.NamedType('responses', univ.SequenceOf(SingleResponse())),
        namedtype.OptionalNamedType('responseExtensions', x509.Extensions().subtype(explicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatSimple, 1)))
        )

class BasicOCSPResponse(univ.Sequence):
    componentType = namedtype.NamedTypes(
        namedtype.NamedType('tbsResponseData', ResponseData()),
        namedtype.NamedType('signatureAlgorithm', x509.AlgorithmIdentifier()),
        namedtype.NamedType('signature', univ.BitString()),
        namedtype.OptionalNamedType('certs', univ.SequenceOf(x509.Certificate()).subtype(explicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatSimple, 0)))
        )

class ResponseBytes(univ.Sequence):
    componentType = namedtype.NamedTypes(
        namedtype.NamedType('responseType', univ.ObjectIdentifier()),
        namedtype.NamedType('response', univ.OctetString())
        )

class OCSPResponseStatus(univ.Enumerated):
    namedValues = namedval.NamedValues(
        ('successful', 0),
        ('malformedRequest', 1),
        ('internalError', 2),
        ('tryLater', 3),
        ('undefinedStatus', 4),  # should never occur
        ('sigRequired', 5),
        ('unauthorized', 6)
        )

class OCSPResponse(univ.Sequence):
    componentType = namedtype.NamedTypes(
        namedtype.NamedType('responseStatus',  OCSPResponseStatus()),
        namedtype.OptionalNamedType('responseBytes', ResponseBytes().subtype(explicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatSimple, 0)))
        )

class Request(univ.Sequence):
    componentType = namedtype.NamedTypes(
        namedtype.NamedType('reqCert', CertID()),
        namedtype.OptionalNamedType('singleRequestExtensions', x509.Extensions().subtype(explicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatSimple, 0)))
        )

class Signature(univ.Sequence):
    componentType = namedtype.NamedTypes(
        namedtype.NamedType('signatureAlgorithm', x509.AlgorithmIdentifier()),
        namedtype.NamedType('signature', univ.BitString()),
        namedtype.OptionalNamedType('certs', univ.SequenceOf(x509.Certificate()).subtype(explicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatSimple, 0)))
        )

class TBSRequest(univ.Sequence):
    componentType = namedtype.NamedTypes(
        namedtype.DefaultedNamedType('version', Version('v1').subtype(explicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatSimple, 0))),
        namedtype.OptionalNamedType('requestorName', GeneralName().subtype(explicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatSimple, 1))),
        namedtype.NamedType('requestList', univ.SequenceOf(Request())),
        namedtype.OptionalNamedType('requestExtensions', x509.Extensions().subtype(explicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatSimple, 2)))
        )

class OCSPRequest(univ.Sequence):
    componentType = namedtype.NamedTypes(
        namedtype.NamedType('tbsRequest',  TBSRequest()),
        namedtype.OptionalNamedType('optionalSignature', Signature().subtype(explicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatSimple, 0)))
        )

# end of OCSP module definitions

sha1oid = univ.ObjectIdentifier((1, 3, 14, 3, 2, 26))

class ValueOnlyBitStringEncoder(encoder.encoder.BitStringEncoder):
    # These methods just do not encode tag and length fields of TLV
    def encodeTag(self, *args): return ''
    def encodeLength(self, *args): return ''
    def encodeValue(*args):
        substrate, isConstructed = encoder.encoder.BitStringEncoder.encodeValue(*args)
        # OCSP-specific hack follows: cut off the "unused bit count"
        # encoded bit-string value.
        return substrate[1:], isConstructed

    def __call__(self, bitStringValue):
        return self.encode(None, bitStringValue, defMode=1, maxChunkSize=0)
       
valueOnlyBitStringEncoder = ValueOnlyBitStringEncoder()

def mkOcspRequest(issuerCert, userCert):
    issuerTbsCertificate = issuerCert.getComponentByName('tbsCertificate')
    issuerSubject = issuerTbsCertificate.getComponentByName('subject')
    
    userTbsCertificate = userCert.getComponentByName('tbsCertificate')
    userIssuer = userTbsCertificate.getComponentByName('issuer')

    assert issuerSubject == userIssuer

    userIssuerHash = hashlib.sha1(
        encoder.encode(userIssuer)
        ).digest()
    
    issuerSubjectPublicKey = issuerTbsCertificate.getComponentByName('subjectPublicKeyInfo').getComponentByName('subjectPublicKey')
    
    issuerKeyHash =  hashlib.sha1(
        valueOnlyBitStringEncoder(issuerSubjectPublicKey)
        ).digest()
    
    userSerialNumber = userTbsCertificate.getComponentByName('serialNumber')

    # Build request object

    request = Request()
    
    reqCert = request.setComponentByName('reqCert').getComponentByName('reqCert')
    
    hashAlgorithm = reqCert.setComponentByName('hashAlgorithm').getComponentByName('hashAlgorithm')
    hashAlgorithm.setComponentByName('algorithm', sha1oid)
    
    reqCert.setComponentByName('issuerNameHash', userIssuerHash)
    reqCert.setComponentByName('issuerKeyHash', issuerKeyHash)
    reqCert.setComponentByName('serialNumber', userSerialNumber)

    ocspRequest = OCSPRequest()
    
    tbsRequest = ocspRequest.setComponentByName('tbsRequest').getComponentByName('tbsRequest')
    tbsRequest.setComponentByName('version', 'v1')
    
    requestList = tbsRequest.setComponentByName('requestList').getComponentByName('requestList')
    requestList.setComponentByPosition(0, request)
    
    return ocspRequest

def parseOcspResponse(ocspResponse):
    responseStatus = ocspResponse.getComponentByName('responseStatus')
    assert responseStatus == OCSPResponseStatus('successful'), responseStatus.prettyPrint()    
    responseBytes = ocspResponse.getComponentByName('responseBytes')
    responseType = responseBytes.getComponentByName('responseType')
    assert responseType == id_pkix_ocsp_basic, responseType.prettyPrint()
    
    response = responseBytes.getComponentByName('response')

    basicOCSPResponse, _ = decoder.decode(
        response, asn1Spec=BasicOCSPResponse()
        )
    
    tbsResponseData = basicOCSPResponse.getComponentByName('tbsResponseData')

    response0 = tbsResponseData.getComponentByName('responses').getComponentByPosition(0)
    
    return (
        tbsResponseData.getComponentByName('producedAt'),
        response0.getComponentByName('certID'),
        response0.getComponentByName('certStatus').getName(),
        response0.getComponentByName('thisUpdate')
        )

# Sample OCSP client
 
if __name__ == '__main__':
    import sys
    import urllib2
    
    if len(sys.argv) != 2:
        print """Usage:
$ cat CACertificate.pem userCertificate.pem | %s <ocsp-responder-url>""" % sys.argv[0]
        sys.exit(-1)
    else:
        ocspUrl = sys.argv[1]

    # Parse CA and user certificates
    
    issuerCert, _ = decoder.decode(
        x509.readPemFromFile(sys.stdin),
        asn1Spec=x509.Certificate()
        )
    userCert, _ = decoder.decode(
        x509.readPemFromFile(sys.stdin),
        asn1Spec=x509.Certificate()
        )

    # Build OCSP request
    
    ocspReq = mkOcspRequest(issuerCert, userCert)

    # Use HTTP POST to get response (see Appendix A of RFC 2560)
    # In case you need proxies, set the http_proxy env variable

    httpReq = urllib2.Request(
        ocspUrl,
        encoder.encode(ocspReq),
        { 'Content-Type': 'application/ocsp-request' }
        )
    httpRsp = urllib2.urlopen(httpReq).read()

    # Process OCSP response
    
    ocspRsp, _ = decoder.decode(httpRsp, asn1Spec=OCSPResponse())

    producedAt, certId, certStatus, thisUpdate = parseOcspResponse(ocspRsp)

    print 'Certificate ID %s is %s at %s till %s\n' % (
        certId.getComponentByName('serialNumber'),
        certStatus,
        producedAt,
        thisUpdate
        )
