<s> from __future__ import unicode_literals <EOL> import os <EOL> import logbook <EOL> import sys <EOL> import trollius as asyncio <EOL> from trollius import From , subprocess , Return <EOL> from trollius . executor import TimeoutError <EOL> from trollius . py33_exceptions import ProcessLookupError <EOL> from pyshark . tshark . tshark import get_tshark_path , get_tshark_display_filter_flag <EOL> from pyshark . tshark . tshark_xml import packet_from_xml_packet , psml_structure_from_xml <EOL> class TSharkCrashException ( Exception ) : <EOL> pass <EOL> class UnknownEncyptionStandardException ( Exception ) : <EOL> pass <EOL> class Capture ( object ) : <EOL> """<STR_LIT>""" <EOL> DEFAULT_BATCH_SIZE = <NUM_LIT:2> ** <NUM_LIT:16> <EOL> SUMMARIES_BATCH_SIZE = <NUM_LIT:64> <EOL> DEFAULT_LOG_LEVEL = logbook . CRITICAL <EOL> SUPPORTED_ENCRYPTION_STANDARDS = [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] <EOL> def __init__ ( self , display_filter = None , only_summaries = False , eventloop = None , <EOL> decryption_key = None , encryption_type = '<STR_LIT>' , output_file = None , <EOL> decode_as = None , tshark_path = None , override_prefs = None ) : <EOL> self . _packets = [ ] <EOL> self . current_packet = <NUM_LIT:0> <EOL> self . display_filter = display_filter <EOL> self . only_summaries = only_summaries <EOL> self . output_file = output_file <EOL> self . running_processes = set ( ) <EOL> self . loaded = False <EOL> self . decode_as = decode_as <EOL> self . log = logbook . Logger ( self . __class__ . __name__ , level = self . DEFAULT_LOG_LEVEL ) <EOL> self . tshark_path = tshark_path <EOL> self . override_prefs = override_prefs <EOL> self . debug = False <EOL> self . eventloop = eventloop <EOL> if self . eventloop is None : <EOL> self . setup_eventloop ( ) <EOL> if encryption_type and encryption_type . lower ( ) in self . SUPPORTED_ENCRYPTION_STANDARDS : <EOL> self . encryption = ( decryption_key , encryption_type . lower ( ) ) <EOL> else : <EOL> raise UnknownEncyptionStandardException ( "<STR_LIT>" <EOL> % '<STR_LIT:U+002CU+0020>' . join ( self . SUPPORTED_ENCRYPTION_STANDARDS ) ) <EOL> def __getitem__ ( self , item ) : <EOL> """<STR_LIT>""" <EOL> return self . _packets [ item ] <EOL> def __len__ ( self ) : <EOL> return len ( self . _packets ) <EOL> def next ( self ) : <EOL> return self . next_packet ( ) <EOL> def next_packet ( self ) : <EOL> if self . current_packet >= len ( self . _packets ) : <EOL> raise StopIteration ( ) <EOL> cur_packet = self . _packets [ self . current_packet ] <EOL> self . current_packet += <NUM_LIT:1> <EOL> return cur_packet <EOL> def clear ( self ) : <EOL> """<STR_LIT>""" <EOL> self . _packets = [ ] <EOL> self . current_packet = <NUM_LIT:0> <EOL> def reset ( self ) : <EOL> """<STR_LIT>""" <EOL> self . current_packet = <NUM_LIT:0> <EOL> def load_packets ( self , packet_count = <NUM_LIT:0> , timeout = None ) : <EOL> """<STR_LIT>""" <EOL> initial_packet_amount = len ( self . _packets ) <EOL> def keep_packet ( pkt ) : <EOL> self . _packets . append ( pkt ) <EOL> if packet_count != <NUM_LIT:0> and len ( self . _packets ) - initial_packet_amount >= packet_count : <EOL> raise Return ( ) <EOL> try : <EOL> self . apply_on_packets ( keep_packet , timeout = timeout ) <EOL> self . loaded = True <EOL> except TimeoutError : <EOL> pass <EOL> def set_debug ( self ) : <EOL> """<STR_LIT>""" <EOL> self . log . level = logbook . DEBUG <EOL> self . debug = True <EOL> def setup_eventloop ( self ) : <EOL> """<STR_LIT>""" <EOL> if os . name == '<STR_LIT>' : <EOL> self . eventloop = asyncio . ProactorEventLoop ( ) <EOL> if sys . version_info <= ( <NUM_LIT:3> , <NUM_LIT:0> ) : <EOL> self . eventloop . _selector . _stopped_serving = set ( ) <EOL> else : <EOL> self . eventloop = asyncio . new_event_loop ( ) <EOL> asyncio . set_event_loop ( self . eventloop ) <EOL> @ staticmethod <EOL> def _extract_tag_from_data ( data , tag_name = b'<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> opening_tag = b'<STR_LIT:<>' + tag_name + b'<STR_LIT:>>' <EOL> closing_tag = opening_tag . replace ( b'<STR_LIT:<>' , b'<STR_LIT>' ) <EOL> tag_end = data . find ( closing_tag ) <EOL> if tag_end != - <NUM_LIT:1> : <EOL> tag_end += len ( closing_tag ) <EOL> tag_start = data . find ( opening_tag ) <EOL> return data [ tag_start : tag_end ] , data [ tag_end : ] <EOL> return None , data <EOL> def _packets_from_tshark_sync ( self , packet_count = None , existing_process = None ) : <EOL> """<STR_LIT>""" <EOL> tshark_process = existing_process or self . eventloop . run_until_complete ( self . _get_tshark_process ( ) ) <EOL> psml_structure , data = self . eventloop . run_until_complete ( self . _get_psml_struct ( tshark_process . stdout ) ) <EOL> packets_captured = <NUM_LIT:0> <EOL> data = b'<STR_LIT>' <EOL> try : <EOL> while True : <EOL> try : <EOL> packet , data = self . eventloop . run_until_complete ( <EOL> self . _get_packet_from_stream ( tshark_process . stdout , data , psml_structure = psml_structure ) ) <EOL> except EOFError : <EOL> self . log . debug ( '<STR_LIT>' ) <EOL> break <EOL> if packet : <EOL> packets_captured += <NUM_LIT:1> <EOL> yield packet <EOL> if packet_count and packets_captured >= packet_count : <EOL> break <EOL> finally : <EOL> self . _cleanup_subprocess ( tshark_process ) <EOL> def apply_on_packets ( self , callback , timeout = None ) : <EOL> """<STR_LIT>""" <EOL> coro = self . packets_from_tshark ( callback ) <EOL> if timeout is not None : <EOL> coro = asyncio . wait_for ( coro , timeout ) <EOL> return self . eventloop . run_until_complete ( coro ) <EOL> @ asyncio . coroutine <EOL> def packets_from_tshark ( self , packet_callback , packet_count = None ) : <EOL> """<STR_LIT>""" <EOL> tshark_process = yield From ( self . _get_tshark_process ( packet_count = packet_count ) ) <EOL> try : <EOL> yield From ( self . _go_through_packets_from_fd ( tshark_process . stdout , packet_callback , <EOL> packet_count = packet_count ) ) <EOL> finally : <EOL> self . _cleanup_subprocess ( tshark_process ) <EOL> @ asyncio . coroutine <EOL> def _go_through_packets_from_fd ( self , fd , packet_callback , packet_count = None ) : <EOL> """<STR_LIT>""" <EOL> packets_captured = <NUM_LIT:0> <EOL> self . log . debug ( '<STR_LIT>' ) <EOL> psml_struct , data = yield From ( self . _get_psml_struct ( fd ) ) <EOL> while True : <EOL> try : <EOL> packet , data = yield From ( self . _get_packet_from_stream ( fd , data , psml_structure = psml_struct ) ) <EOL> except EOFError : <EOL> self . log . debug ( '<STR_LIT>' ) <EOL> break <EOL> if packet : <EOL> packets_captured += <NUM_LIT:1> <EOL> packet_callback ( packet ) <EOL> if packet_count and packets_captured >= packet_count : <EOL> break <EOL> @ asyncio . coroutine <EOL> def _get_psml_struct ( self , fd ) : <EOL> """<STR_LIT>""" <EOL> data = b'<STR_LIT>' <EOL> psml_struct = None <EOL> if self . only_summaries : <EOL> while not psml_struct : <EOL> new_data = yield From ( fd . read ( self . SUMMARIES_BATCH_SIZE ) ) <EOL> data += new_data <EOL> psml_struct , data = self . _extract_tag_from_data ( data , b'<STR_LIT>' ) <EOL> if psml_struct : <EOL> psml_struct = psml_structure_from_xml ( psml_struct ) <EOL> elif not new_data : <EOL> raise Return ( None , data ) <EOL> raise Return ( psml_struct , data ) <EOL> else : <EOL> raise Return ( None , data ) <EOL> @ asyncio . coroutine <EOL> def _get_packet_from_stream ( self , stream , existing_data , psml_structure = None ) : <EOL> """<STR_LIT>""" <EOL> packet , existing_data = self . _extract_tag_from_data ( existing_data ) <EOL> if packet : <EOL> packet = packet_from_xml_packet ( packet , psml_structure = psml_structure ) <EOL> raise Return ( packet , existing_data ) <EOL> new_data = yield From ( stream . read ( self . DEFAULT_BATCH_SIZE ) ) <EOL> existing_data += new_data <EOL> if not new_data : <EOL> raise EOFError ( ) <EOL> raise Return ( None , existing_data ) <EOL> @ asyncio . coroutine <EOL> def _get_tshark_process ( self , packet_count = None , stdin = None ) : <EOL> """<STR_LIT>""" <EOL> xml_type = '<STR_LIT>' if self . only_summaries else '<STR_LIT>' <EOL> parameters = [ get_tshark_path ( self . tshark_path ) , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , xml_type ] + self . get_parameters ( packet_count = packet_count ) <EOL> self . log . debug ( '<STR_LIT>' + '<STR_LIT:U+0020>' . join ( parameters ) ) <EOL> output = None if self . debug else open ( os . devnull , "<STR_LIT:w>" ) <EOL> tshark_process = yield From ( asyncio . create_subprocess_exec ( * parameters , <EOL> stdout = subprocess . PIPE , <EOL> stderr = output , <EOL> stdin = stdin ) ) <EOL> self . log . debug ( '<STR_LIT>' ) <EOL> if tshark_process . returncode is not None and tshark_process . returncode != <NUM_LIT:0> : <EOL> raise TSharkCrashException ( <EOL> '<STR_LIT>' % '<STR_LIT:U+0020>' . join ( parameters ) ) <EOL> self . running_processes . add ( tshark_process ) <EOL> raise Return ( tshark_process ) <EOL> def _cleanup_subprocess ( self , process ) : <EOL> """<STR_LIT>""" <EOL> if process . returncode is None : <EOL> try : <EOL> process . kill ( ) <EOL> except ProcessLookupError : <EOL> pass <EOL> except OSError : <EOL> if os . name != '<STR_LIT>' : <EOL> raise <EOL> elif process . returncode > <NUM_LIT:0> : <EOL> raise TSharkCrashException ( '<STR_LIT>' % process . returncode ) <EOL> def close ( self ) : <EOL> for process in self . running_processes : <EOL> self . _cleanup_subprocess ( process ) <EOL> def __del__ ( self ) : <EOL> self . close ( ) <EOL> def get_parameters ( self , packet_count = None ) : <EOL> """<STR_LIT>""" <EOL> params = [ ] <EOL> if self . display_filter : <EOL> params += [ get_tshark_display_filter_flag ( self . tshark_path ) , self . display_filter ] <EOL> if packet_count : <EOL> params += [ '<STR_LIT:-c>' , str ( packet_count ) ] <EOL> if all ( self . encryption ) : <EOL> params += [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' + self . encryption [ <NUM_LIT:1> ] + '<STR_LIT>' + <EOL> self . encryption [ <NUM_LIT:0> ] + '<STR_LIT:">' ] <EOL> if self . override_prefs : <EOL> for preference_name , preference_value in self . override_prefs . items ( ) : <EOL> if all ( self . encryption ) and preference_name in ( '<STR_LIT>' , '<STR_LIT>' ) : <EOL> continue <EOL> params += [ '<STR_LIT>' , '<STR_LIT>' . format ( preference_name , preference_value ) ] <EOL> if self . output_file : <EOL> params += [ '<STR_LIT>' , self . output_file ] <EOL> if self . decode_as : <EOL> for criterion , decode_as_proto in self . decode_as . items ( ) : <EOL> params += [ '<STR_LIT>' , '<STR_LIT:U+002C>' . join ( [ criterion . strip ( ) , decode_as_proto . strip ( ) ] ) ] <EOL> return params <EOL> def __iter__ ( self ) : <EOL> if self . loaded : <EOL> return iter ( self . _packets ) <EOL> else : <EOL> return self . _packets_from_tshark_sync ( ) <EOL> def __repr__ ( self ) : <EOL> return '<STR_LIT>' % ( self . __class__ . __name__ , len ( self . _packets ) ) </s>
