From 58a9d337d921d4ea3e96de390bd803acdf44f738 Mon Sep 17 00:00:00 2001 From: cuteNBD <206630735+bingdongni@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:43:39 +0800 Subject: [PATCH] Preserve time zone offsets in session values --- lib/core/convert.py | 85 +++++++++++++++++++++++-- tests/test_serialize.py | 136 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 217 insertions(+), 4 deletions(-) diff --git a/lib/core/convert.py b/lib/core/convert.py index f3927e60d64..805de33640d 100644 --- a/lib/core/convert.py +++ b/lib/core/convert.py @@ -61,6 +61,83 @@ "lib.utils.har.RawPair", )) +class _SessionTimezone(datetime.tzinfo): + """Fixed timezone information captured from a stored date or time value.""" + + def __init__(self, offset, daylight, name): + self._offset = offset + self._daylight = daylight + self._name = name + + def utcoffset(self, dt): + return self._offset + + def dst(self, dt): + return self._daylight + + def tzname(self, dt): + return self._name + + def __reduce__(self): + return self.__class__, (self._offset, self._daylight, self._name) + + def fromutc(self, dt): + if not isinstance(dt, datetime.datetime): + raise TypeError("fromutc() requires a datetime argument") + if dt.tzinfo is not self: + raise ValueError("fromutc: dt.tzinfo is not self") + if self._offset is None: + raise ValueError("fromutc: non-None utcoffset() result required") + return dt + self._offset + +def _serializeTimeDelta(value): + return [value.days, value.seconds, value.microseconds] if value is not None else None + +def _serializeDateTime(value, retVal): + if value.tzinfo is not None: + name = value.tzname() + if isinstance(name, six.binary_type): + name = {"b": encodeBase64(name, binary=False)} + retVal["z"] = [_serializeTimeDelta(value.utcoffset()), _serializeTimeDelta(value.dst()), name] + if getattr(value, "fold", 0): + retVal["f"] = value.fold + return retVal + +def _deserializeTimeDelta(value): + if value is None: + return None + if not isinstance(value, list) or len(value) != 3 or not all(isinstance(_, six.integer_types) for _ in value): + raise ValueError("invalid serialized timezone offset") + retVal = datetime.timedelta(*value) + if not -datetime.timedelta(days=1) < retVal < datetime.timedelta(days=1): + raise ValueError("serialized timezone offset must be less than one day") + return retVal + +def _deserializeDateTime(value, struct): + kwargs = {} + if "z" in struct: + zone = struct["z"] + if not isinstance(zone, list) or len(zone) != 3: + raise ValueError("invalid serialized timezone") + name = zone[2] + if isinstance(name, dict): + if set(name) != set(("b",)) or not isinstance(name["b"], six.string_types): + raise ValueError("invalid serialized timezone name") + name = decodeBase64(name["b"], binary=True) + if six.PY3: + try: + name = name.decode("utf-8") + except UnicodeDecodeError: + name = name.decode("latin-1") + elif name is not None and not isinstance(name, six.string_types): + raise ValueError("invalid serialized timezone name") + elif six.PY2 and isinstance(name, six.text_type): + name = name.encode("utf-8") + kwargs["tzinfo"] = _SessionTimezone(_deserializeTimeDelta(zone[0]), _deserializeTimeDelta(zone[1]), name) + if getattr(value, "fold", None) is not None and "f" in struct: + kwargs["fold"] = struct["f"] + return value.replace(**kwargs) if kwargs else value + def _serializeEncode(value): """ Turns a Python value into a JSON-serializable (tagged) structure @@ -126,13 +203,13 @@ def _serializeEncode(value): return {_SERIALIZE_TAG: "dec", "v": getUnicode(value)} if isinstance(value, datetime.datetime): - return {_SERIALIZE_TAG: "dt", "v": [value.year, value.month, value.day, value.hour, value.minute, value.second, value.microsecond]} + return _serializeDateTime(value, {_SERIALIZE_TAG: "dt", "v": [value.year, value.month, value.day, value.hour, value.minute, value.second, value.microsecond]}) if isinstance(value, datetime.date): return {_SERIALIZE_TAG: "date", "v": [value.year, value.month, value.day]} if isinstance(value, datetime.time): - return {_SERIALIZE_TAG: "time", "v": [value.hour, value.minute, value.second, value.microsecond]} + return _serializeDateTime(value, {_SERIALIZE_TAG: "time", "v": [value.hour, value.minute, value.second, value.microsecond]}) if isinstance(value, datetime.timedelta): return {_SERIALIZE_TAG: "td", "v": [value.days, value.seconds, value.microseconds]} @@ -206,11 +283,11 @@ def _serializeDecode(struct): elif tag == "dec": return decimal.Decimal(struct["v"]) elif tag == "dt": - return datetime.datetime(*struct["v"]) + return _deserializeDateTime(datetime.datetime(*struct["v"]), struct) elif tag == "date": return datetime.date(*struct["v"]) elif tag == "time": - return datetime.time(*struct["v"]) + return _deserializeDateTime(datetime.time(*struct["v"]), struct) elif tag == "td": return datetime.timedelta(struct["v"][0], struct["v"][1], struct["v"][2]) elif tag == "o": diff --git a/tests/test_serialize.py b/tests/test_serialize.py index fb7b72bcb71..5b656ff7481 100644 --- a/tests/test_serialize.py +++ b/tests/test_serialize.py @@ -25,6 +25,7 @@ Kept deliberately verbose - each case maps to a real session/BigArray value or a real report. """ +import copy import datetime import decimal import json @@ -53,6 +54,22 @@ def rt(value): return deserializeValue(serializeValue(value)) +class _TestTimezone(datetime.tzinfo): + def __init__(self, offset, dst, name): + self.offset = offset + self.daylight = dst + self.name = name + + def utcoffset(self, dt): + return self.offset + + def dst(self, dt): + return self.daylight + + def tzname(self, dt): + return self.name + + def _import_all_modules(): """Import every sqlmap module (like --smoke-test) so class-subclass reflection sees them all.""" root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) @@ -259,6 +276,102 @@ def test_datetime_family(self): self.assertEqual(restored, value) self.assertIs(type(restored), type(value)) + def test_timezone_aware_datetime_and_time(self): + zones = [ + _TestTimezone(datetime.timedelta(0), None, "UTC"), + _TestTimezone(datetime.timedelta(hours=5, minutes=30), datetime.timedelta(0), "IST"), + _TestTimezone(datetime.timedelta(hours=-7), datetime.timedelta(hours=1), "PDT"), + _TestTimezone(None, None, None), + ] + for zone in zones: + for value in (datetime.datetime(2024, 1, 2, 3, 4, 5, 6, tzinfo=zone), + datetime.time(3, 4, 5, 6, tzinfo=zone)): + restored = rt(value) + self.assertEqual(restored, value) + self.assertIs(type(restored), type(value)) + self.assertIsNotNone(restored.tzinfo) + self.assertEqual(restored.utcoffset(), value.utcoffset()) + self.assertEqual(restored.dst(), value.dst()) + self.assertEqual(restored.tzname(), value.tzname()) + self.assertEqual(restored.isoformat(), value.isoformat()) + + @unittest.skipUnless(hasattr(datetime.datetime.now(), "fold"), "requires datetime.fold") + def test_datetime_fold_roundtrip(self): + zone = _TestTimezone(datetime.timedelta(hours=-5), datetime.timedelta(0), "EST") + for value in (datetime.datetime(2024, 11, 3, 1, 30, tzinfo=zone).replace(fold=1), + datetime.time(1, 30, tzinfo=zone).replace(fold=1), + datetime.datetime(2024, 11, 3, 1, 30).replace(fold=1)): + restored = rt(value) + self.assertEqual(restored.fold, 1) + self.assertEqual(restored.isoformat(), value.isoformat()) + + def test_legacy_datetime_serialization(self): + self.assertEqual(deserializeValue('{"$T":"dt","v":[2024,1,2,3,4,5,6]}'), + datetime.datetime(2024, 1, 2, 3, 4, 5, 6)) + self.assertEqual(deserializeValue('{"$T":"time","v":[3,4,5,6]}'), datetime.time(3, 4, 5, 6)) + + def test_restored_timezone_converts_from_utc(self): + zone = _TestTimezone(datetime.timedelta(hours=5, minutes=30), None, "IST") + value = datetime.datetime(2024, 1, 2, 3, 4, tzinfo=zone) + restored = rt(value) + utc = _TestTimezone(datetime.timedelta(0), datetime.timedelta(0), "UTC") + instant = datetime.datetime(2024, 1, 1, 21, 34, tzinfo=utc) + self.assertEqual(instant.astimezone(restored.tzinfo), value) + self.assertEqual(copy.deepcopy(restored), restored) + + def test_timezone_name_with_non_ascii_characters(self): + name = u"Caf\u00e9" + if six.PY2: + name = name.encode("utf-8") + zone = _TestTimezone(datetime.timedelta(0), None, name) + value = datetime.datetime(2024, 1, 2, tzinfo=zone) + self.assertEqual(rt(value).tzname(), name) + + @unittest.skipUnless(six.PY2, "Python 2 permits non-UTF-8 byte timezone names") + def test_timezone_name_with_non_utf8_bytes(self): + name = b"Zone\xff" + zone = _TestTimezone(datetime.timedelta(0), None, name) + value = datetime.datetime(2024, 1, 2, tzinfo=zone) + self.assertEqual(rt(value).tzname(), name) + + def test_timezone_name_from_python2_session(self): + payload = '{"$T":"dt","v":[2024,1,2,3,4,5,6],"z":[[0,0,0],null,{"b":"Wm9uZf8="}]}' + expected = b"Zone\xff" if six.PY2 else u"Zone\u00ff" + self.assertEqual(deserializeValue(payload).tzname(), expected) + + @unittest.skipUnless(sys.version_info >= (3, 7), "requires sub-minute timezone offsets") + def test_timezone_with_subminute_offset(self): + zone = _TestTimezone(datetime.timedelta(hours=5, seconds=42, microseconds=123), None, "TST") + value = datetime.datetime(2024, 1, 2, tzinfo=zone) + restored = rt(value) + self.assertEqual(restored.utcoffset(), value.utcoffset()) + self.assertEqual(restored.isoformat(), value.isoformat()) + + @unittest.skipUnless(hasattr(datetime.datetime.now(), "fold"), "requires datetime.fold") + def test_ambiguous_datetime_keeps_its_instant(self): + class FoldTimezone(datetime.tzinfo): + def utcoffset(self, dt): + return datetime.timedelta(hours=-5 if dt.fold else -4) + + def dst(self, dt): + return datetime.timedelta(0) if dt.fold else datetime.timedelta(hours=1) + + def tzname(self, dt): + return "EST" if dt.fold else "EDT" + + value = datetime.datetime(2024, 11, 3, 1, 30, tzinfo=FoldTimezone()).replace(fold=1) + restored = rt(value) + self.assertEqual(restored.fold, value.fold) + self.assertEqual(restored.isoformat(), value.isoformat()) + self.assertEqual(restored.timestamp(), value.timestamp()) + + def test_invalid_serialized_timezone_is_rejected(self): + for zone in ("invalid", [[], None, "UTC"], [[1, 0, 0], None, "UTC"], [None, None, 42], + [None, None, {"unknown": "value"}]): + payload = json.dumps({"$T": "dt", "v": [2024, 1, 2, 3, 4, 5, 6], "z": zone}) + with self.assertRaises(ValueError): + deserializeValue(payload) + class TestRealSessionObjects(unittest.TestCase): # the exact shapes written with serialize=True across the codebase @@ -422,6 +535,29 @@ def test_known_serializable_classes_present(self): class TestHashDBIntegration(unittest.TestCase): + def test_timezone_survives_session_reload(self): + from lib.utils.hashdb import HashDB + + handle, path = tempfile.mkstemp(suffix=".sqlite") + os.close(handle) + os.remove(path) + try: + value = datetime.datetime(2024, 1, 2, 3, 4, tzinfo=_TestTimezone( + datetime.timedelta(hours=5, minutes=30), datetime.timedelta(0), "IST")) + db = HashDB(path) + db.write("query", [(value,)], serialize=True) + db.flush() + db._write_cache.clear() + db._read_cache.cache.clear() + restored = db.retrieve("query", unserialize=True) + db.closeAll() + + self.assertEqual(getUnicode(restored), getUnicode([(value,)])) + self.assertEqual(restored[0][0].utcoffset(), value.utcoffset()) + finally: + if os.path.exists(path): + os.remove(path) + def test_write_retrieve_roundtrip(self): from lib.utils.hashdb import HashDB