raise exception when kasten loads invalid packed bytes

This commit is contained in:
Kevin Froman 2020-12-24 05:00:36 +00:00
parent 903e203c63
commit 67be9077ab
2 changed files with 22 additions and 4 deletions

View File

@ -14,3 +14,7 @@ class InvalidID(KastenException):
"""Occurs when a Kasten generator fails validation, which means data does not match KastenChecksum""" """Occurs when a Kasten generator fails validation, which means data does not match KastenChecksum"""
pass pass
class InvalidPackedBytes(KastenException):
"""Occurs when packed bytes cannot be deserialized"""
pass

View File

@ -1,9 +1,13 @@
"""Main Kasten object, does nothing but provide access to packed Kasten bytes and call a specified generator function""" """Main Kasten object, does nothing but provide access to packed Kasten bytes and call a specified generator function"""
import traceback
from msgpack import unpackb from msgpack import unpackb
import msgpack
from .types import KastenChecksum from .types import KastenChecksum
from .types import KastenPacked from .types import KastenPacked
from .generator import pack from .generator import pack
from .exceptions import InvalidPackedBytes
""" """
Copyright (C) <2020> Kevin Froman Copyright (C) <2020> Kevin Froman
@ -31,10 +35,20 @@ class Kasten:
**additional_generator_kwargs): # noqa **additional_generator_kwargs): # noqa
self.id = id self.id = id
self.generator = generator self.generator = generator
header, data = packed_bytes.split(b'\n', 1) try:
header = unpackb( header, data = packed_bytes.split(b'\n', 1)
header, except ValueError:
strict_map_key=True) raise InvalidPackedBytes("Could not extract data section")
try:
header = unpackb(
header,
strict_map_key=True)
except( # noqa
msgpack.exceptions.ExtraData,
msgpack.FormatError,
UnicodeEncodeError,
ValueError) as _:
raise InvalidPackedBytes("Could not decode packed bytes")
self.header = header self.header = header
self.data = data self.data = data
self.additional_generator_args = list(additional_generator_args) self.additional_generator_args = list(additional_generator_args)