def stream(self):
"""
Yields Flow objects from the dump.
"""
# There is a weird mingw bug that breaks .tell() when reading from stdin.
try:
self.fo.tell()
except IOError: # pragma: no cover
can_tell = False
else:
can_tell = True
off = 0
try:
while True:
data = tnetstring.load(self.fo)
try:
data = io_compat.migrate_flow(data)
except ValueError as e:
raise exceptions.FlowReadException(str(e))
if can_tell:
off = self.fo.tell()
if data["type"] not in models.FLOW_TYPES:
raise exceptions.FlowReadException(
"Unknown flow type: {}".format(data["type"]))
yield models.FLOW_TYPES[data["type"]].from_state(data)
except ValueError:
# Error is due to EOF
if can_tell and self.fo.tell() == off and self.fo.read() == b'':
return
raise exceptions.FlowReadException("Invalid data format.")
async def load_flows_from_path(self, path: str) -> int:
path = os.path.expanduser(path)
try:
with open(path, "rb") as f:
return await self.load_flows(f)
except IOError as e:
ctx.log.error("Cannot load flows: {}".format(e))
raise exceptions.FlowReadException(str(e)) from e
def stream(self) -> Iterable[flow.Flow]:
"""
Yields Flow objects from the dump.
"""
try:
while True:
data = tnetstring.load(self.fo)
try:
data = io_compat.migrate_flow(data)
except ValueError as e:
raise exceptions.FlowReadException(str(e))
if data["type"] not in FLOW_TYPES:
raise exceptions.FlowReadException("Unknown flow type: {}".format(data["type"]))
yield FLOW_TYPES[data["type"]].from_state(data)
except ValueError as e:
if str(e) == "not a tnetstring: empty file":
return # Error is due to EOF
raise exceptions.FlowReadException("Invalid data format.")
def load_flows_file(self, path):
path = os.path.expanduser(path)
try:
if path == "-":
# This is incompatible with Python 3 - maybe we can use click?
freader = io.FlowReader(sys.stdin)
return self.load_flows(freader)
else:
with open(path, "rb") as f:
freader = io.FlowReader(f)
return self.load_flows(freader)
except IOError as v:
raise exceptions.FlowReadException(v.strerror)
def stream(self) -> Iterable[flow.Flow]:
"""
Yields Flow objects from the dump.
"""
try:
while True:
# FIXME: This cast hides a lack of dynamic type checking
loaded = cast(
Dict[Union[bytes, str], Any],
tnetstring.load(self.fo),
)
try:
mdata = compat.migrate_flow(loaded)
except ValueError as e:
raise exceptions.FlowReadException(str(e))
if mdata["type"] not in FLOW_TYPES:
raise exceptions.FlowReadException(
"Unknown flow type: {}".format(mdata["type"]))
yield FLOW_TYPES[mdata["type"]].from_state(mdata)
except (ValueError, TypeError, IndexError) as e:
if str(e) == "not a tnetstring: empty file":
return # Error is due to EOF
raise exceptions.FlowReadException("Invalid data format.")
def load_flows(self, fo: typing.IO[bytes]) -> int:
cnt = 0
freader = io.FlowReader(fo)
try:
for flow in freader.stream():
ctx.master.load_flow(flow)
cnt += 1
except (IOError, exceptions.FlowReadException) as e:
if cnt:
ctx.log.warn("Flow file corrupted - loaded %i flows." % cnt)
else:
ctx.log.error("Flow file corrupted.")
raise exceptions.FlowReadException(str(e)) from e
else:
return cnt
def load_flows_file(self, path: str) -> int:
path = os.path.expanduser(path)
try:
if path == "-":
try:
sys.stdin.buffer.read(0)
except Exception as e:
raise IOError("Cannot read from stdin: {}".format(e))
freader = io.FlowReader(sys.stdin.buffer)
return self.load_flows(freader)
else:
with open(path, "rb") as f:
freader = io.FlowReader(f)
return self.load_flows(freader)
except IOError as v:
raise exceptions.FlowReadException(v.strerror)
def load_flows_file(self, path: str) -> int:
path = os.path.expanduser(path)
cnt = 0
try:
with open(path, "rb") as f:
freader = io.FlowReader(f)
for i in freader.stream():
cnt += 1
ctx.master.load_flow(i)
return cnt
except (IOError, exceptions.FlowReadException) as v:
if cnt:
ctx.log.warn("Flow file corrupted - loaded %i flows." % cnt, )
else:
ctx.log.error("Flow file corrupted.")
raise exceptions.FlowReadException(v)
def read_flows_from_paths(paths):
"""
Given a list of filepaths, read all flows and return a list of them.
From a performance perspective, streaming would be advisable -
however, if there's an error with one of the files, we want it to be raised immediately.
Raises:
FlowReadException, if any error occurs.
"""
try:
flows = []
for path in paths:
path = os.path.expanduser(path)
with open(path, "rb") as f:
flows.extend(FlowReader(f).stream())
except IOError as e:
raise exceptions.FlowReadException(e.strerror)
return flows
