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)