def __iter__(self):
def read_chunk_header(ammo_file):
chunk_header = ''
while chunk_header is '':
line = ammo_file.readline()
if line.startswith('['):
self.headers.update(_parse_header(line.strip('\r\n[]\t ')))
elif line is '':
return line
else:
chunk_header = line.strip('\r\n')
return chunk_header
opener = get_opener(self.filename)
with opener(self.filename, 'rb') as ammo_file:
info.status.af_size = opener.data_length
# if we got StopIteration here, the file is empty
chunk_header = read_chunk_header(ammo_file)
while chunk_header:
if chunk_header is not '':
try:
fields = chunk_header.split()
chunk_size = int(fields[0])
if chunk_size == 0:
self.log.debug(
'Zero-sized chunk in ammo file at %s. Starting over.'
% ammo_file.tell())
ammo_file.seek(0)
info.status.inc_loop_count()
chunk_header = read_chunk_header(ammo_file)
continue
uri = fields[1]
marker = fields[2] if len(fields) > 2 else None
missile = ammo_file.read(chunk_size)
if len(missile) < chunk_size:
raise AmmoFileError(
"Unexpected end of file: read %s bytes instead of %s"
% (len(missile), chunk_size))
yield (HttpAmmo(
uri=uri,
headers=[
': '.join(header)
for header in self.headers.items()
],
method='POST',
body=missile,
http_ver=self.http_ver,
).to_s(), marker)
except (IndexError, ValueError) as e:
raise AmmoFileError(
"Error while reading ammo file. Position: %s, header: '%s', original exception: %s"
% (ammo_file.tell(), chunk_header, e))
chunk_header = read_chunk_header(ammo_file)
if chunk_header == '':
self.log.debug(
'Reached the end of ammo file. Starting over.')
ammo_file.seek(0)
info.status.inc_loop_count()
chunk_header = read_chunk_header(ammo_file)
info.status.af_position = ammo_file.tell()
def __iter__(self):
opener = get_opener(self.filename)
with opener(self.filename, 'rb') as ammo_file:
info.status.af_size = opener.data_length
while True:
for line in ammo_file:
info.status.af_position = ammo_file.tell()
if line.startswith('['):
self.headers.add(line.strip('\r\n[]\t '))
elif len(line.rstrip('\r\n')):
fields = line.split()
uri = fields[0]
if len(fields) > 1:
marker = fields[1]
else:
marker = None
yield (
HttpAmmo(
uri,
headers=self.headers,
http_ver=self.http_ver,
).to_s(), marker)
if info.status.ammo_count == 0:
self.log.error("No ammo in uri-style file")
raise AmmoFileError("No ammo! Cover me!")
ammo_file.seek(0)
info.status.af_position = 0
info.status.inc_loop_count()
def __iter__(self):
def read_chunk_header(ammo_file):
chunk_header = ''
while chunk_header is '':
line = ammo_file.readline()
if line is '':
return line
chunk_header = line.strip('\r\n')
return chunk_header
opener = get_opener(self.filename)
with opener(self.filename, 'rb') as ammo_file:
info.status.af_size = opener.data_length
# if we got StopIteration here, the file is empty
chunk_header = read_chunk_header(ammo_file)
while chunk_header:
if chunk_header is not '':
try:
fields = chunk_header.split()
chunk_size = int(fields[0])
if chunk_size == 0:
if info.status.loop_count == 0:
self.log.info(
'Zero-sized chunk in ammo file at %s. Starting over.'
% ammo_file.tell())
ammo_file.seek(0)
info.status.inc_loop_count()
chunk_header = read_chunk_header(ammo_file)
continue
marker = fields[1] if len(fields) > 1 else None
missile = ammo_file.read(chunk_size)
if len(missile) < chunk_size:
raise AmmoFileError(
"Unexpected end of file: read %s bytes instead of %s"
% (len(missile), chunk_size))
yield (missile, marker)
except (IndexError, ValueError) as e:
raise AmmoFileError(
"Error while reading ammo file. Position: %s, header: '%s', original exception: %s"
% (ammo_file.tell(), chunk_header, e))
chunk_header = read_chunk_header(ammo_file)
if chunk_header == '':
ammo_file.seek(0)
info.status.inc_loop_count()
chunk_header = read_chunk_header(ammo_file)
info.status.af_position = ammo_file.tell()
def get_ammo_generator(self):
"""
return ammo generator
"""
af_readers = {
'phantom': missile.AmmoFileReader,
'slowlog': missile.SlowLogReader,
'line': missile.LineReader,
'uri': missile.UriReader,
'uripost': missile.UriPostReader,
'access': missile.AccessLogReader,
}
if self.uris and self.ammo_file:
raise StepperConfigurationError(
'Both uris and ammo file specified. You must specify only one of them'
)
elif self.uris:
ammo_gen = missile.UriStyleGenerator(self.uris,
self.headers,
http_ver=self.http_ver)
elif self.ammo_file:
if self.ammo_type in af_readers:
if self.ammo_type is 'phantom':
with get_opener(self.ammo_file)(self.ammo_file,
'rb') as ammo:
try:
if not ammo.next()[0].isdigit():
self.ammo_type = 'uri'
self.log.info(
"Setting ammo_type 'uri' because ammo is not started with digit and you did not specify ammo format"
)
else:
self.log.info(
"Default ammo type ('phantom') used, use 'phantom.ammo_type' option to override it"
)
except StopIteration, e:
self.log(
"Couldn't read first line of ammo file: %s" %
e)
raise AmmoFileError(
"Couldn't read first line of ammo file")
else:
raise NotImplementedError(
'No such ammo type implemented: "%s"' % self.ammo_type)
ammo_gen = af_readers[self.ammo_type](self.ammo_file,
headers=self.headers,
http_ver=self.http_ver)