<commit_before>from __future__ import with_statement
from contextlib import contextmanager
from Network import TCPServer, TCPClient, UDPServer, UDPClient, _NamedCache

from templates import Protocol, UInt, PDU, MessageTemplate, Char, StructTemplate, ListTemplate, UnionTemplate
from binary_conversions import to_0xhex, to_bin


class Rammbock(object):

    def __init__(self):
        self._init_caches()

    def _init_caches(self):
        self._protocol_in_progress = None
        self._protocols = {}
        self._servers = _NamedCache('server')
        self._clients = _NamedCache('client')
        self._message_stack = []
        self._default_values = None

    def reset_rammbock(self):
        """Closes all connections, deletes all servers, clients, and protocols.

        You should call this method before exiting your test run. This will close all the connections and the ports
        will therefore be available for reuse faster.
        """
        for server in self._servers:
            server.close()
        for client in self._clients:
            client.close()
        self._init_caches()

    def start_protocol_description(self, protocol_name):
        """Start defining a new protocol template.

        All messages sent and received from a connection that uses a protocol have to conform to this protocol template.
        Protocol template fields can be used to search messages from buffer.
        """
        if self._protocol_in_progress:
            raise Exception('Can not start a new protocol definition in middle of old.')
        if protocol_name in self._protocols:
            raise Exception('Protocol %s already defined' % protocol_name)
        self._protocol_in_progress = Protocol(protocol_name)

    def end_protocol_description(self):
        """End protocol definition."""
        self._protocols[self._protocol_in_progress.name] = self._protocol_in_progress
        self._protocol_in_progress = None

    def start_udp_server(self, ip, port, name=None, timeout=None, protocol=None):
        protocol = self._get_protocol(protocol)
        server = UDPServer(ip=ip, port=port, timeout=timeout, protocol=protocol)
        return self._servers.add(server, name)

    def start_udp_client(self, ip=None, port=None, name=None, timeout=None, protocol=None):
        self._start_client(UDPClient, ip, port, name, timeout, protocol)

    def start_tcp_client(self, ip=None, port=None, name=None, timeout=None, protocol=None):
        self._start_client(TCPClient, ip, port, name, timeout, protocol)

    def _start_client(self, client_class, ip=None, port=None, name=None, timeout=None, protocol=None):
        protocol = self._get_protocol(protocol)
        client = client_class(timeout=timeout, protocol=protocol)
        if ip or port:
            client.set_own_ip_and_port(ip=ip, port=port)
        return self._clients.add(client, name)

    def _get_protocol(self, protocol):
        protocol = self._protocols[protocol] if protocol else None
        return protocol

    def start_tcp_server(self, ip, port, name=None, timeout=None, protocol=None):
        protocol = self._get_protocol(protocol)
        server = TCPServer(ip=ip, port=port, timeout=timeout, protocol=protocol)
        return self._servers.add(server, name)

    def get_client_protocol(self, name):
        return self._clients.get(name).protocol

    def accept_connection(self, name=None, alias=None):
        server = self._servers.get(name)
        server.accept_connection(alias)

    def connect(self, host, port, name=None):
        """Connect a client to certain host and port."""
        client = self._clients.get(name)
        client.connect_to(host, port)

    # TODO: Log the raw binary that is sent and received.
    def client_sends_binary(self, message, name=None):
        """Send raw binary data."""<commit_msg>refactor start server to callback helper<commit_after>from __future__ import with_statement
from contextlib import contextmanager
from Network import TCPServer, TCPClient, UDPServer, UDPClient, _NamedCache

from templates import Protocol, UInt, PDU, MessageTemplate, Char, StructTemplate, ListTemplate, UnionTemplate
from binary_conversions import to_0xhex, to_bin


class Rammbock(object):

    def __init__(self):
        self._init_caches()

    def _init_caches(self):
        self._protocol_in_progress = None
        self._protocols = {}
        self._servers = _NamedCache('server')
        self._clients = _NamedCache('client')
        self._message_stack = []
        self._default_values = None

    def reset_rammbock(self):
        """Closes all connections, deletes all servers, clients, and protocols.

        You should call this method before exiting your test run. This will close all the connections and the ports
        will therefore be available for reuse faster.
        """
        for server in self._servers:
            server.close()
        for client in self._clients:
            client.close()
        self._init_caches()

    def start_protocol_description(self, protocol_name):
        """Start defining a new protocol template.

        All messages sent and received from a connection that uses a protocol have to conform to this protocol template.
        Protocol template fields can be used to search messages from buffer.
        """
        if self._protocol_in_progress:
            raise Exception('Can not start a new protocol definition in middle of old.')
        if protocol_name in self._protocols:
            raise Exception('Protocol %s already defined' % protocol_name)
        self._protocol_in_progress = Protocol(protocol_name)

    def end_protocol_description(self):
        """End protocol definition."""
        self._protocols[self._protocol_in_progress.name] = self._protocol_in_progress
        self._protocol_in_progress = None

    def start_udp_server(self, ip, port, name=None, timeout=None, protocol=None):
        self._start_server(UDPServer, ip, port, name, timeout, protocol)

    def start_tcp_server(self, ip, port, name=None, timeout=None, protocol=None):
        self._start_server(TCPServer, ip, port, name, timeout, protocol)

    def _start_server(self, server_class, ip, port, name=None, timeout=None, protocol=None):
        protocol = self._get_protocol(protocol)
        server = server_class(ip=ip, port=port, timeout=timeout, protocol=protocol)
        return self._servers.add(server, name)

    def start_udp_client(self, ip=None, port=None, name=None, timeout=None, protocol=None):
        self._start_client(UDPClient, ip, port, name, timeout, protocol)

    def start_tcp_client(self, ip=None, port=None, name=None, timeout=None, protocol=None):
        self._start_client(TCPClient, ip, port, name, timeout, protocol)

    def _start_client(self, client_class, ip=None, port=None, name=None, timeout=None, protocol=None):
        protocol = self._get_protocol(protocol)
        client = client_class(timeout=timeout, protocol=protocol)
        if ip or port:
            client.set_own_ip_and_port(ip=ip, port=port)
        return self._clients.add(client, name)

    def _get_protocol(self, protocol):
        protocol = self._protocols[protocol] if protocol else None
        return protocol

    def get_client_protocol(self, name):
        return self._clients.get(name).protocol

    def accept_connection(self, name=None, alias=None):
        server = self._servers.get(name)
        server.accept_connection(alias)

    def connect(self, host, port, name=None):
        """Connect a client to certain host and port."""
        client = self._clients.get(name)
        client.connect_to(host, port)

    # TODO: Log the raw binary that is sent and received.
    def client_sends_binary(self, message, name=None):
        """Send raw binary data."""
{'fraction_non_alphanumeric': 0.08812211390456644, 'fraction_numerical': 0.0002565418163160595, 'mean_word_length': 3.3510044642857144, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6097536', 'n_tokens_mistral': 2139, 'n_tokens_neox': 2109, 'n_words': 636}