<s> import threading <EOL> import subprocess <EOL> import re <EOL> import os <EOL> import shlex <EOL> import time <EOL> import curlbomb <EOL> import sys <EOL> from io import TextIOWrapper , BytesIO <EOL> from tempfile import NamedTemporaryFile <EOL> import logging <EOL> import socket <EOL> import unittest <EOL> import tornado <EOL> import tempfile <EOL> import hashlib <EOL> import shutil <EOL> log = logging . getLogger ( "<STR_LIT>" ) <EOL> log . setLevel ( level = logging . INFO ) <EOL> client_scripts = { <EOL> '<STR_LIT>' : ( "<STR_LIT>" , "<STR_LIT>" ) , <EOL> '<STR_LIT>' : ( "<STR_LIT>" , "<STR_LIT>" ) , <EOL> '<STR_LIT>' : ( "<STR_LIT>" , "<STR_LIT>" ) , <EOL> '<STR_LIT>' : ( "<STR_LIT>" , "<STR_LIT>" ) <EOL> } <EOL> class CurlbombThread ( threading . Thread ) : <EOL> def __init__ ( self , settings ) : <EOL> threading . Thread . __init__ ( self ) <EOL> self . settings = settings <EOL> self . returncode = None <EOL> def run ( self ) : <EOL> log . info ( '<STR_LIT>' ) <EOL> try : <EOL> self . returncode = curlbomb . run_server ( self . settings ) <EOL> finally : <EOL> log . info ( '<STR_LIT>' ) <EOL> class CurlbombTestBase ( unittest . TestCase ) : <EOL> def get_curlbomb ( self , args , script = None ) : <EOL> """<STR_LIT>""" <EOL> if type ( script ) == str : <EOL> script = bytes ( script , "<STR_LIT:utf-8>" ) <EOL> stdin = "<STR_LIT>" not in args and script is not None <EOL> try : <EOL> override_defaults = { } <EOL> log . info ( "<STR_LIT>" . format ( stdin ) ) <EOL> if stdin : <EOL> s = TextIOWrapper ( BytesIO ( script ) ) <EOL> override_defaults [ '<STR_LIT>' ] = s <EOL> else : <EOL> s = NamedTemporaryFile ( ) <EOL> if script is not None : <EOL> s . write ( script ) <EOL> s . flush ( ) <EOL> args = args . format ( script = s . name ) <EOL> args = shlex . split ( args ) <EOL> log . warn ( "<STR_LIT>" . format ( args ) ) <EOL> settings = curlbomb . get_settings ( args , override_defaults ) <EOL> client_cmd = settings [ '<STR_LIT>' ] ( settings ) <EOL> curlbomb_thread = CurlbombThread ( settings ) <EOL> curlbomb_thread . start ( ) <EOL> return ( curlbomb_thread , <EOL> client_cmd ) <EOL> finally : <EOL> s . close ( ) <EOL> def run_client ( self , client_cmd , expected_out = None , expected_err = None ) : <EOL> client_cmd = [ '<STR_LIT>' , '<STR_LIT:-c>' , client_cmd ] <EOL> log . info ( "<STR_LIT>" . format ( client_cmd ) ) <EOL> client_proc = subprocess . Popen ( client_cmd , stdout = subprocess . PIPE , <EOL> stderr = subprocess . PIPE ) <EOL> client_out , client_err = client_proc . communicate ( ) <EOL> client_out = client_out . decode ( "<STR_LIT:utf-8>" ) <EOL> client_err = client_err . decode ( "<STR_LIT:utf-8>" ) <EOL> log . info ( '<STR_LIT>' . format ( repr ( client_out ) ) ) <EOL> if expected_out is not None : <EOL> self . assertEquals ( client_out , expected_out ) <EOL> if expected_err is not None : <EOL> self . assertIsNotNone ( expected_err . search ( client_err ) ) <EOL> return client_out , client_err <EOL> def simple_runner ( self , args , script , expected_out ) : <EOL> cb , client_cmd = self . get_curlbomb ( args , script ) <EOL> client_out , client_err = self . run_client ( client_cmd , expected_out ) <EOL> return ( cb , client_cmd , client_out , client_err ) <EOL> def test_default_args ( self ) : <EOL> self . simple_runner ( '<STR_LIT>' , * client_scripts [ '<STR_LIT>' ] ) <EOL> def test_no_knock ( self ) : <EOL> self . simple_runner ( '<STR_LIT>' , * client_scripts [ '<STR_LIT>' ] ) <EOL> def test_knock ( self ) : <EOL> script , expected_out = client_scripts [ '<STR_LIT>' ] <EOL> cb , client_cmd = self . get_curlbomb ( '<STR_LIT>' , script ) <EOL> self . assertTrue ( client_cmd . startswith ( "<STR_LIT>" ) ) <EOL> client_cmd_no_knock = "<STR_LIT:U+0020>" . join ( client_cmd . split ( "<STR_LIT:U+0020>" ) [ <NUM_LIT:1> : ] ) <EOL> client_out , client_err = self . run_client ( client_cmd_no_knock , '<STR_LIT>' ) <EOL> client_out , client_err = self . run_client ( client_cmd , expected_out ) <EOL> def test_multi_gets ( self ) : <EOL> script , expected_out = client_scripts [ '<STR_LIT>' ] <EOL> cb , client_cmd = self . get_curlbomb ( '<STR_LIT>' , script ) <EOL> for x in range ( <NUM_LIT:4> ) : <EOL> self . run_client ( client_cmd , expected_out ) <EOL> time . sleep ( <NUM_LIT:1> ) <EOL> self . run_client ( client_cmd , '<STR_LIT>' , re . compile ( "<STR_LIT>" ) ) <EOL> def test_specific_port ( self ) : <EOL> s = socket . socket ( socket . AF_INET , socket . SOCK_STREAM ) <EOL> s . bind ( ( '<STR_LIT>' , <NUM_LIT:0> ) ) <EOL> port = s . getsockname ( ) [ <NUM_LIT:1> ] <EOL> s . close ( ) <EOL> self . simple_runner ( '<STR_LIT>' . format ( port ) , * client_scripts [ '<STR_LIT>' ] ) <EOL> def test_python ( self ) : <EOL> self . simple_runner ( '<STR_LIT>' , * client_scripts [ '<STR_LIT>' ] ) <EOL> def test_alternate_command ( self ) : <EOL> self . simple_runner ( '<STR_LIT>' , * client_scripts [ '<STR_LIT>' ] ) <EOL> def test_wget ( self ) : <EOL> self . simple_runner ( '<STR_LIT>' , * client_scripts [ '<STR_LIT>' ] ) <EOL> def test_survey ( self ) : <EOL> self . simple_runner ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ) <EOL> def test_survey_wget ( self ) : <EOL> self . simple_runner ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ) <EOL> def test_unwrapped_command ( self ) : <EOL> script , expected_out = client_scripts [ '<STR_LIT>' ] <EOL> cb , client_cmd = self . get_curlbomb ( '<STR_LIT>' , script ) <EOL> self . assertTrue ( client_cmd . startswith ( '<STR_LIT>' ) ) <EOL> client_out , client_err = self . run_client ( client_cmd , expected_out ) <EOL> def test_domain ( self ) : <EOL> s = socket . socket ( socket . AF_INET , socket . SOCK_STREAM ) <EOL> s . bind ( ( '<STR_LIT>' , <NUM_LIT:0> ) ) <EOL> port = s . getsockname ( ) [ <NUM_LIT:1> ] <EOL> s . close ( ) <EOL> cb , client_cmd , out , err = self . simple_runner ( <EOL> '<STR_LIT>' . format ( port = port ) , <EOL> * client_scripts [ '<STR_LIT>' ] ) <EOL> self . assertTrue ( "<STR_LIT>" . format ( port = port ) in client_cmd ) <EOL> def test_log_posts ( self ) : <EOL> script , expected_out = client_scripts [ '<STR_LIT>' ] <EOL> cb , client_cmd = self . get_curlbomb ( '<STR_LIT>' , script ) <EOL> try : <EOL> original_stdout = sys . stdout <EOL> out = sys . stdout = TextIOWrapper ( BytesIO ( ) ) <EOL> client_out , client_err = self . run_client ( client_cmd , expected_out ) <EOL> out . seek ( <NUM_LIT:0> ) <EOL> self . assertEquals ( out . read ( ) , client_out ) <EOL> finally : <EOL> sys . stdout = original_stdout <EOL> def test_unwrapped ( self ) : <EOL> self . simple_runner ( '<STR_LIT>' , * client_scripts [ '<STR_LIT>' ] ) <EOL> def __get_directory_contents ( self , path ) : <EOL> """<STR_LIT>""" <EOL> parent_path , path_name = os . path . split ( path ) <EOL> contents = [ ] <EOL> for root , dirs , files in os . walk ( path ) : <EOL> contents . append ( os . path . relpath ( root , parent_path ) + '<STR_LIT:/>' ) <EOL> for f in files : <EOL> contents . append ( os . path . relpath ( os . path . join ( root , f ) , parent_path ) ) <EOL> return sorted ( contents ) <EOL> def __get_directory_sha256 ( self , path ) : <EOL> """<STR_LIT>""" <EOL> parent_path , path_name = os . path . split ( path ) <EOL> path_shas = { } <EOL> for root , dirs , files in os . walk ( path ) : <EOL> for fn in ( os . path . join ( root , f ) for f in files ) : <EOL> with open ( fn ) as f : <EOL> path_shas [ os . path . relpath ( fn , parent_path ) ] = hashlib . sha256 ( <EOL> f . read ( ) . encode ( '<STR_LIT:utf-8>' ) ) . hexdigest ( ) <EOL> return path_shas <EOL> def __put_get_test ( self , operation ) : <EOL> curdir = os . path . abspath ( os . path . curdir ) <EOL> try : <EOL> test_path = os . path . join ( <EOL> os . path . dirname ( os . path . abspath ( __file__ ) ) , <EOL> '<STR_LIT>' , '<STR_LIT>' ) <EOL> contents = self . __get_directory_contents ( test_path ) <EOL> contents . remove ( os . path . join ( <EOL> '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ) ) <EOL> test_shas = self . __get_directory_sha256 ( test_path ) <EOL> del test_shas [ os . path . join ( <EOL> '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ) ] <EOL> def dest_assert ( client_out , destdir ) : <EOL> if operation == '<STR_LIT>' : <EOL> self . assertEqual ( sorted ( client_out . splitlines ( ) ) , contents ) <EOL> dest_contents = self . __get_directory_contents ( <EOL> os . path . join ( destdir , '<STR_LIT>' ) ) <EOL> self . assertEqual ( contents , dest_contents ) <EOL> self . assertEqual ( test_shas , self . __get_directory_sha256 ( <EOL> os . path . join ( destdir , '<STR_LIT>' ) ) ) <EOL> with tempfile . TemporaryDirectory ( suffix = '<STR_LIT>' ) as tempdir : <EOL> os . chdir ( tempdir ) <EOL> with tempfile . TemporaryDirectory ( ) as destdir : <EOL> try : <EOL> cb , client_cmd = self . get_curlbomb ( <EOL> '<STR_LIT>' . format ( <EOL> operation = operation , source = test_path , dest = destdir ) ) <EOL> client_out , client_err = self . run_client ( client_cmd ) <EOL> cb . join ( ) <EOL> dest_assert ( client_out , destdir ) <EOL> finally : <EOL> pass <EOL> with tempfile . TemporaryDirectory ( suffix = '<STR_LIT>' ) as tmpdir : <EOL> os . chdir ( tmpdir ) <EOL> cb , client_cmd = self . get_curlbomb ( <EOL> '<STR_LIT>' . format ( <EOL> operation = operation , source = test_path ) ) <EOL> client_out , client_err = self . run_client ( client_cmd ) <EOL> cb . join ( ) <EOL> dest_assert ( client_out , tmpdir ) <EOL> finally : <EOL> os . chdir ( curdir ) <EOL> def test_put ( self ) : <EOL> self . __put_get_test ( '<STR_LIT>' ) <EOL> def test_get ( self ) : <EOL> self . __put_get_test ( '<STR_LIT>' ) <EOL> def test_ssl ( self ) : <EOL> script , expected_out = client_scripts [ '<STR_LIT>' ] <EOL> ca_dir = os . path . join ( os . path . dirname ( os . path . realpath ( __file__ ) ) , <EOL> '<STR_LIT>' ) <EOL> ca_pem = os . path . join ( ca_dir , '<STR_LIT>' ) <EOL> ca_chain = os . path . join ( ca_dir , '<STR_LIT>' ) <EOL> cb , client_cmd = self . get_curlbomb ( <EOL> '<STR_LIT>' . format ( ca_pem = ca_pem ) , script ) <EOL> client_out , client_err = self . run_client ( client_cmd ) <EOL> print ( client_err ) <EOL> self . assertTrue ( "<STR_LIT>" in client_err ) <EOL> os . environ [ '<STR_LIT>' ] = ca_chain <EOL> try : <EOL> client_out , client_err = self . run_client ( client_cmd , expected_out ) <EOL> finally : <EOL> del os . environ [ '<STR_LIT>' ] <EOL> def test_ssl_self_signed ( self ) : <EOL> """<STR_LIT>""" <EOL> self . simple_runner ( '<STR_LIT>' , * client_scripts [ '<STR_LIT>' ] ) <EOL> def test_ping ( self ) : <EOL> args = '<STR_LIT>' <EOL> cb , client_cmd = self . get_curlbomb ( args ) <EOL> self . assertTrue ( client_cmd . startswith ( '<STR_LIT>' ) ) <EOL> self . assertIn ( '<STR_LIT>' . format ( cb . settings [ '<STR_LIT>' ] ) , client_cmd ) <EOL> client_cmd_no_knock = re . sub ( '<STR_LIT>' , '<STR_LIT>' , <EOL> client_cmd ) <EOL> client_out , client_err = self . run_client ( client_cmd_no_knock ) <EOL> self . assertIn ( "<STR_LIT>" , client_out ) <EOL> client_out , client_err = self . run_client ( client_cmd ) <EOL> cb . join ( ) <EOL> self . assertEquals ( cb . returncode , <NUM_LIT:0> ) <EOL> def test_multi_ping_return ( self ) : <EOL> """<STR_LIT>""" <EOL> args = '<STR_LIT>' <EOL> cb , client_cmd = self . get_curlbomb ( args ) <EOL> for x in range ( <NUM_LIT:3> ) : <EOL> self . run_client ( client_cmd ) <EOL> cb . join ( ) <EOL> self . assertEquals ( cb . returncode , <NUM_LIT:0> ) <EOL> args = '<STR_LIT>' <EOL> cb , client_cmd = self . get_curlbomb ( args ) <EOL> self . run_client ( client_cmd ) <EOL> self . run_client ( re . sub ( '<STR_LIT>' , '<STR_LIT>' , client_cmd ) ) <EOL> self . run_client ( client_cmd ) <EOL> cb . join ( ) <EOL> self . assertEquals ( cb . returncode , <NUM_LIT> ) </s>
