taxii-json.py

import uuid

# http://marshmallow.readthedocs.org/en/latest/
from marshmallow import Schema, fields

STATUS_TYPES = [
    'ASYNCHRONOUS_POLL_ERROR', 'BAD_MESSAGE', 'DENIED',
    'DESTINATION_COLLECTION_ERROR', 'FAILURE', 'INVALID_RESPONSE_PART',
    'NETWORK_ERROR', 'NOT_FOUND', 'PENDING', 'POLLING_UNSUPPORTED',
    'RETRY', 'SUCCESS', 'UNAUTHORIZED', 'UNSUPPORTED_MESSAGE',
    'UNSUPPORTED_CONTENT', 'UNSUPPORTED_PROTOCOL', 'UNSUPPORTED_QUERY'
]

STATUS_DETAIL_FIELD = [
    'ACCEPTABLE_DESTINATION', 'MAX_PART_NUMBER', 'ITEM',
    'ESTIMATED_WAIT', 'RESULT_ID', 'WILL_PUSH', 'ESTIMATED_WAIT',
    'SUPPORTED_BINDING', 'SUPPORTED_CONTENT',
    'SUPPORTED_PROTOCOL', 'SUPPORTED_PROTOCOL'
]

SERVICE_TYPES = ['DISCOVERY', 'COLLECTION_MANAGEMENT', 'INBOX', 'POLL']


# Schemes

class TAXIIMessageSchema(Schema):

    # old message_id. Forcing UUID usage here
    id = fields.UUID(required=True)

    extended_headers = fields.Raw()
    signature = fields.String(default=None)


class TAXIIResponseSchema(TAXIIMessageSchema):

    # forcing UUID usage here
    in_response_to = fields.UUID(required=True)


class StatusMessageSchema(TAXIIResponseSchema):

    # old status_type
    type = fields.Enum(choices=STATUS_TYPES, required=True)

    message = fields.String()
    extended_headers = fields.Raw()

    # all status_details
    details = fields.Raw()


class DiscoveryRequestSchema(TAXIIMessageSchema):
    pass


class ContentBindingsSchema(Schema):

    # old binding_id
    id = fields.String(required=True)

    # old subtype_ids
    subtypes = fields.List(fields.String)


class QuerySchema(Schema):

    format_id = fields.String(required=True)
    value = fields.Raw(required=True)


class ServiceInstanceSchema(Schema):

    address = fields.URL()

    # old service_type
    type = fields.Enum(choices=SERVICE_TYPES, required=True)

    # old service_version
    version = fields.String()

    available = fields.Boolean()

    message_bindings = fields.List(fields.String)
    supported_queries = fields.Nested(QuerySchema, many=True)

    content_bindings = fields.Nested(ContentBindingsSchema, many=True)
    message = fields.String()



class DiscoveryResponseSchema(TAXIIResponseSchema):

    extended_headers = fields.Raw()

    # old service_instances
    services = fields.Nested(ServiceInstanceSchema, many=True)


# Entitites

class DiscoveryRequest(object):

    def __init__(self, id, signature=None, extended_headers=None):
        self.id = id
        self.signature = signature
        self.extended_headers = extended_headers


# JSON serialization

discovery_request = dict(id=uuid.uuid4(),
        extended_headers={'custom-header' : 'header-value'})

result = DiscoveryRequestSchema().dumps(discovery_request)
print result.data
# {"extended_headers": {"custom-header": "header-value"}, "id": "73c62c2b-c8f8-44dc-bb28-922363a69da2", "signature": null}


# Object serialization

discovery_request = DiscoveryRequest(id=uuid.uuid4(), signature="signature-value")

result = DiscoveryRequestSchema().dumps(discovery_request)
print result.data
# {"extended_headers": null, "id": "312fd904-8632-4b61-bf05-9ad80c4a3aa3", "signature": "signature-value"}