<commit_before>from __future__ import unicode_literals
import os
import logbook
import sys

import trollius as asyncio
from logbook import StreamHandler
from trollius import From, subprocess, Return
from trollius.executor import TimeoutError
from trollius.py33_exceptions import ProcessLookupError

from pyshark.tshark.tshark import get_process_path, get_tshark_display_filter_flag, \
    tshark_supports_json, TSharkVersionException
from pyshark.tshark.tshark_json import packet_from_json_packet
from pyshark.tshark.tshark_xml import packet_from_xml_packet, psml_structure_from_xml


class TSharkCrashException(Exception):
    pass


class UnknownEncyptionStandardException(Exception):
    pass


class StopCapture(Exception):
    """
    Exception that the user can throw anywhere in packet-handling to stop the capture process.
    """
    pass


class Capture(object):
    """
    Base class for packet captures.
    """
    DEFAULT_BATCH_SIZE = 2 ** 16
    SUMMARIES_BATCH_SIZE = 64
    DEFAULT_LOG_LEVEL = logbook.CRITICAL
    SUPPORTED_ENCRYPTION_STANDARDS = ['wep', 'wpa-pwk', 'wpa-pwd', 'wpa-psk']

    def __init__(self, display_filter=None, only_summaries=False, eventloop=None,
                 decryption_key=None, encryption_type='wpa-pwd', output_file=None,
                 decode_as=None,  disable_protocol=None, tshark_path=None,
                 override_prefs=None, capture_filter=None, use_json=False):

        self.loaded = False
        self.tshark_path = tshark_path
        self._override_prefs = override_prefs
        self.debug = False
        self.use_json = use_json
        self._packets = []
        self._current_packet = 0
        self._display_filter = display_filter
        self._capture_filter = capture_filter
        self._only_summaries = only_summaries
        self._output_file = output_file
        self._running_processes = set()
        self._decode_as = decode_as
        self._disable_protocol = disable_protocol
        self._log = logbook.Logger(self.__class__.__name__, level=self.DEFAULT_LOG_LEVEL)

        self.eventloop = eventloop
        if self.eventloop is None:
            self._setup_eventloop()
        if encryption_type and encryption_type.lower() in self.SUPPORTED_ENCRYPTION_STANDARDS:
            self.encryption = (decryption_key, encryption_type.lower())
        else:
            raise UnknownEncyptionStandardException("Only the following standards are supported: %s."
                                                    % ', '.join(self.SUPPORTED_ENCRYPTION_STANDARDS))

    def __getitem__(self, item):
        """
        Gets the packet in the given index.

        :param item: packet index
        :return: Packet object.
        """
        return self._packets[item]

    def __len__(self):
        return len(self._packets)

    def next(self):
        return self.next_packet()

    # Allows for child classes to call next() from super() without 2to3 "fixing"
    # the call
    def next_packet(self):
        if self._current_packet >= len(self._packets):
            raise StopIteration()
        cur_packet = self._packets[self._current_packet]
        self._current_packet += 1
        return cur_packet

    def clear(self):
        """
        Empties the capture of any saved packets.
        """
        self._packets = []
        self._current_packet = 0

    def reset(self):
        """
        Starts iterating packets from the first one.
        """
        self._current_packet = 0

    def load_packets(self, packet_count=0, timeout=None):
        """
        Reads the packets from the source (cap, interface, etc.) and adds it to the internal list.
        If 0 as the packet_count is given, reads forever

        :param packet_count: The amount of packets to add to the packet list (0 to read forever)
        :param timeout: If given, automatically stops after a given amount of time.
        """
        initial_packet_amount = len(self._packets)
        def keep_packet(pkt):
            self._packets.append(pkt)

            if packet_count != 0 and len(self._packets) - initial_packet_amount >= packet_count:
                raise Return()

        try:
            self.apply_on_packets(keep_packet, timeout=timeout)
            self.loaded = True
        except TimeoutError:
            pass

    def set_debug(self, set_to=True):
        """
        Sets the capture to debug mode (or turns it off if specified).
        """
        if set_to:
            StreamHandler(sys.stdout).push_application()
            self._log.level = logbook.DEBUG
        self.debug = set_to

    def _setup_eventloop(self):
        """
        Sets up a new eventloop as the current one according to the OS.
        """
        if os.name == 'nt':
            self.eventloop = asyncio.ProactorEventLoop()
            if sys.version_info <= (3, 0):
                # FIXME: There appears to be a bug in the 2.7 version of trollius, wherein the selector retrieves an
                # object of value 0 and attempts to look for it in the weakref set, which raises an exception.
                # This hack sidesteps this issue, but does not solve it. If a proper fix is found, apply it!
                self.eventloop._selector._stopped_serving = set()
        else:
            self.eventloop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.eventloop)
        asyncio.get_child_watcher().attach_loop(self.eventloop)

    @classmethod
    def _get_json_separator(cls):
        return ("}%s%s  ," % (os.linesep, os.linesep)).encode()

    @classmethod
    def _extract_packet_json_from_data(cls, data, got_first_packet=True):
        tag_start = 0
        if not got_first_packet:
            tag_start = data.find(b"{")
            if tag_start == -1:
                return None, data
        closing_tag = cls._get_json_separator()
        tag_end = data.find(closing_tag)
        if tag_end == -1:
            closing_tag = ("}%s%s]" % (os.linesep, os.linesep)).encode()
            tag_end = data.find(closing_tag)
        if tag_end != -1:
            # Include closing parenthesis but not comma<commit_msg>Only applying the fix on unix.<commit_after>from __future__ import unicode_literals
import os
import logbook
import sys

import trollius as asyncio
from logbook import StreamHandler
from trollius import From, subprocess, Return
from trollius.executor import TimeoutError
from trollius.py33_exceptions import ProcessLookupError

from pyshark.tshark.tshark import get_process_path, get_tshark_display_filter_flag, \
    tshark_supports_json, TSharkVersionException
from pyshark.tshark.tshark_json import packet_from_json_packet
from pyshark.tshark.tshark_xml import packet_from_xml_packet, psml_structure_from_xml


class TSharkCrashException(Exception):
    pass


class UnknownEncyptionStandardException(Exception):
    pass


class StopCapture(Exception):
    """
    Exception that the user can throw anywhere in packet-handling to stop the capture process.
    """
    pass


class Capture(object):
    """
    Base class for packet captures.
    """
    DEFAULT_BATCH_SIZE = 2 ** 16
    SUMMARIES_BATCH_SIZE = 64
    DEFAULT_LOG_LEVEL = logbook.CRITICAL
    SUPPORTED_ENCRYPTION_STANDARDS = ['wep', 'wpa-pwk', 'wpa-pwd', 'wpa-psk']

    def __init__(self, display_filter=None, only_summaries=False, eventloop=None,
                 decryption_key=None, encryption_type='wpa-pwd', output_file=None,
                 decode_as=None,  disable_protocol=None, tshark_path=None,
                 override_prefs=None, capture_filter=None, use_json=False):

        self.loaded = False
        self.tshark_path = tshark_path
        self._override_prefs = override_prefs
        self.debug = False
        self.use_json = use_json
        self._packets = []
        self._current_packet = 0
        self._display_filter = display_filter
        self._capture_filter = capture_filter
        self._only_summaries = only_summaries
        self._output_file = output_file
        self._running_processes = set()
        self._decode_as = decode_as
        self._disable_protocol = disable_protocol
        self._log = logbook.Logger(self.__class__.__name__, level=self.DEFAULT_LOG_LEVEL)

        self.eventloop = eventloop
        if self.eventloop is None:
            self._setup_eventloop()
        if encryption_type and encryption_type.lower() in self.SUPPORTED_ENCRYPTION_STANDARDS:
            self.encryption = (decryption_key, encryption_type.lower())
        else:
            raise UnknownEncyptionStandardException("Only the following standards are supported: %s."
                                                    % ', '.join(self.SUPPORTED_ENCRYPTION_STANDARDS))

    def __getitem__(self, item):
        """
        Gets the packet in the given index.

        :param item: packet index
        :return: Packet object.
        """
        return self._packets[item]

    def __len__(self):
        return len(self._packets)

    def next(self):
        return self.next_packet()

    # Allows for child classes to call next() from super() without 2to3 "fixing"
    # the call
    def next_packet(self):
        if self._current_packet >= len(self._packets):
            raise StopIteration()
        cur_packet = self._packets[self._current_packet]
        self._current_packet += 1
        return cur_packet

    def clear(self):
        """
        Empties the capture of any saved packets.
        """
        self._packets = []
        self._current_packet = 0

    def reset(self):
        """
        Starts iterating packets from the first one.
        """
        self._current_packet = 0

    def load_packets(self, packet_count=0, timeout=None):
        """
        Reads the packets from the source (cap, interface, etc.) and adds it to the internal list.
        If 0 as the packet_count is given, reads forever

        :param packet_count: The amount of packets to add to the packet list (0 to read forever)
        :param timeout: If given, automatically stops after a given amount of time.
        """
        initial_packet_amount = len(self._packets)
        def keep_packet(pkt):
            self._packets.append(pkt)

            if packet_count != 0 and len(self._packets) - initial_packet_amount >= packet_count:
                raise Return()

        try:
            self.apply_on_packets(keep_packet, timeout=timeout)
            self.loaded = True
        except TimeoutError:
            pass

    def set_debug(self, set_to=True):
        """
        Sets the capture to debug mode (or turns it off if specified).
        """
        if set_to:
            StreamHandler(sys.stdout).push_application()
            self._log.level = logbook.DEBUG
        self.debug = set_to

    def _setup_eventloop(self):
        """
        Sets up a new eventloop as the current one according to the OS.
        """
        if os.name == 'nt':
            self.eventloop = asyncio.ProactorEventLoop()
            if sys.version_info <= (3, 0):
                # FIXME: There appears to be a bug in the 2.7 version of trollius, wherein the selector retrieves an
                # object of value 0 and attempts to look for it in the weakref set, which raises an exception.
                # This hack sidesteps this issue, but does not solve it. If a proper fix is found, apply it!
                self.eventloop._selector._stopped_serving = set()
        else:
            self.eventloop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.eventloop)
        if os.name == 'posix':
            asyncio.get_child_watcher().attach_loop(self.eventloop)

    @classmethod
    def _get_json_separator(cls):
        return ("}%s%s  ," % (os.linesep, os.linesep)).encode()

    @classmethod
    def _extract_packet_json_from_data(cls, data, got_first_packet=True):
        tag_start = 0
        if not got_first_packet:
            tag_start = data.find(b"{")
            if tag_start == -1:
                return None, data
        closing_tag = cls._get_json_separator()
        tag_end = data.find(closing_tag)
        if tag_end == -1:
            closing_tag = ("}%s%s]" % (os.linesep, os.linesep)).encode()
            tag_end = data.find(closing_tag)
        if tag_end != -1:
            # Include closing parenthesis but not comma
{'fraction_non_alphanumeric': 0.08079739625711961, 'fraction_numerical': 0.004231082180634662, 'mean_word_length': 2.417964404894327, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 6}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1339273', 'n_tokens_mistral': 3628, 'n_tokens_neox': 3489, 'n_words': 947}