capture_fuzz_gen.py

"""Convert from transport socket capture to fuzz test case.

Converts from envoy.data.tap.v2alpha.Trace proto to
test.integration.CaptureFuzzTestCase.

Usage: capture_fuzz_gen.py <listener capture> [<cluster capture>]
"""
from __future__ import print_function

import functools
import sys

from google.protobuf import empty_pb2
from google.protobuf import text_format

from envoy.data.tap.v2alpha import capture_pb2
from test.integration import capture_fuzz_pb2


# Collapse adjacent event in the trace that are of the same type.
def Coalesce(trace):
    if not trace.events:
        return []
    events = [trace.events[0]]
    for event in trace.events[1:]:
        if events[-1].HasField('read') and event.HasField('read'):
            events[-1].read.data += event.read.data
        elif events[-1].HasField('write') and event.HasField('write'):
            events[-1].write.data += event.write.data
        else:
            events.append(event)
    return events


# Convert from transport socket Event to test Event.
def ToTestEvent(direction, event):
    test_event = capture_fuzz_pb2.Event()
    if event.HasField('read'):
        setattr(test_event, '%s_send_bytes' % direction, event.read.data)
    elif event.HasField('write'):
        getattr(test_event, '%s_recv_bytes' % direction).MergeFrom(empty_pb2.Empty())
    return test_event


def ToDownstreamTestEvent(event):
    return ToTestEvent('downstream', event)


def ToUpstreamTestEvent(event):
    return ToTestEvent('upstream', event)


# Zip together the listener/cluster events to produce a single trace for replay.
def TestCaseGen(listener_events, cluster_events):
    test_case = capture_fuzz_pb2.CaptureFuzzTestCase()
    if not listener_events:
        return test_case
    test_case.events.extend([ToDownstreamTestEvent(listener_events[0])])
    del listener_events[0]
    while listener_events or cluster_events:
        if not listener_events:
            test_case.events.extend(map(ToUpstreamTestEvent, cluster_events))
            return test_case
        if not cluster_events:
            test_case.events.extend(map(ToDownstreamTestEvent, listener_events))
            return test_case
        if listener_events[0].timestamp.ToDatetime() < cluster_events[0].timestamp.ToDatetime():
            test_case.events.extend([ToDownstreamTestEvent(listener_events[0])])
            del listener_events[0]
        test_case.events.extend([ToUpstreamTestEvent(cluster_events[0])])
        del cluster_events[0]


def CaptureFuzzGen(listener_path, cluster_path=None):
    listener_trace = capture_pb2.Trace()
    with open(listener_path, 'r') as f:
        text_format.Merge(f.read(), listener_trace)
    listener_events = Coalesce(listener_trace)

    cluster_trace = capture_pb2.Trace()
    if cluster_path:
        with open(cluster_path, 'r') as f:
            text_format.Merge(f.read(), cluster_trace)
        cluster_events = Coalesce(cluster_trace)

    print(TestCaseGen(listener_events, cluster_events))


if __name__ == '__main__':
    if len(sys.argv) < 2 or len(sys.argv) > 3:
        print('Usage: %s <listener capture> [<cluster capture>]' % sys.argv[0])
        sys.exit(1)
    CaptureFuzzGen(*sys.argv[1:])