Preserve time zone offsets in session values

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":

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