Refactor pickle imports to use pickletools.optimize

server/src/uds/REST/methods/system.py

@@ -34,6 +34,7 @@ import codecs
 import datetime
 import logging
 import pickle  # nosec: pickle is used to cache data, not to load it
+import pickletools
 import typing
 
 from uds import models
@@ -95,7 +96,7 @@ def get_servicepools_counters(
                 }
                 for x in stats
             ]
-            
+
             # val = [
             #     {
             #         'stamp': x[0],
@@ -114,7 +115,11 @@ def get_servicepools_counters(
 
             # logger.debug('val: %s', val)
             if len(val) >= 2:
-                cache.put(cache_key, codecs.encode(pickle.dumps(val), 'zip'), CACHE_TIME * 2)
+                cache.put(
+                    cache_key,
+                    codecs.encode(pickletools.optimize(pickle.dumps(val, protocol=-1)), 'zip'),
+                    CACHE_TIME * 2,
+                )
             else:
                 val = [{'stamp': since, 'value': 0}, {'stamp': to, 'value': 0}]
         else:

server/src/uds/core/util/storage.py

@@ -34,6 +34,7 @@ import pickle  # nosec: This is e controled pickle use
 import base64
 import hashlib
 import codecs
+import pickletools
 import typing
 import collections.abc
 import logging
@@ -268,8 +269,15 @@ class Storage:
         attr1: typing.Optional[str] = None,
     ) -> None:
         return self.save_to_db(
-            skey, pickle.dumps(data), attr1
+            skey,
-        )  # Protocol 2 is compatible with python 2.7. This will be unnecesary when fully migrated
+            pickletools.optimize(
+                pickle.dumps(
+                    data,
+                    protocol=-1,
+                )
+            ),
+            attr1,
+        )
 
     def update_to_db(
         self,

server/src/uds/models/ticket_store.py

@@ -31,6 +31,7 @@ Author: Adolfo Gómez, dkmaster at dkmon dot com
 import datetime
 import pickle  # nosec: Tickets are generated by us, so we know they are safe
 import logging
+import pickletools
 import typing
 import collections.abc
 
@@ -47,6 +48,7 @@ from .user_service import UserService
 
 logger = logging.getLogger(__name__)
 
+
 class TicketStore(UUIDModel):
     """
     Tickets storing on DB
@@ -76,9 +78,10 @@ class TicketStore(UUIDModel):
 
     @staticmethod
     def generate_uuid() -> str:
-        """In fact, generates a random string of TICKET_LENGTH chars, that will be used as uuid for the ticket (but is not an uuid compliant string)
+        """In fact, generates a random string of TICKET_LENGTH chars, that will be used as uuid for the ticket (but is not an uuid compliant string)"""
-        """
+        return (
-        return CryptoManager().random_string(consts.ticket.TICKET_LENGTH).lower()  # Temporary fix lower() for compat with 3.0
+            CryptoManager().random_string(consts.ticket.TICKET_LENGTH).lower()
+        )  # Temporary fix lower() for compat with 3.0
 
     @staticmethod
     def create(
@@ -98,13 +101,17 @@ class TicketStore(UUIDModel):
         Returns:
             The ticket id
         """
-        data = pickle.dumps(data)
+        data = pickletools.optimize(
+            pickle.dumps(data, protocol=-1)
+        )  # nosec: Tickets are generated by us, so we know they are safe
 
         if secure:
             if not owner:
                 raise ValueError('Tried to use a secure ticket without owner')
             data = CryptoManager().aes_crypt(data, owner.encode())
-            owner = consts.ticket.TICKET_SECURED_ONWER  # So data is REALLY encrypted, because key used to encrypt is sustituted by SECURED on DB
+            owner = (
+                consts.ticket.TICKET_SECURED_ONWER
+            )  # So data is REALLY encrypted, because key used to encrypt is sustituted by SECURED on DB
 
         return TicketStore.objects.create(
             uuid=TicketStore.generate_uuid(),
@@ -169,7 +176,7 @@ class TicketStore(UUIDModel):
                 if not owner:
                     raise ValueError('Tried to use a secure ticket without owner')
                 key = owner.encode()
-            
+
             t = TicketStore.objects.get(uuid=uuid)
 
             data: bytes = t.data
@@ -179,18 +186,20 @@ class TicketStore(UUIDModel):
                     raise ValueError('Tried to use a secure ticket without owner')
                 data = CryptoManager().aes_decrypt(data, key)
 
-            dct = pickle.loads(data)  # nosec: Tickets are ONLY generated by us, so we know they are safe
+            saved_data = pickle.loads(data)  # nosec: Tickets are ONLY generated by us, so we know they are safe
 
             # invoke check function
-            if checkFnc(dct) is False:
+            if checkFnc(saved_data) is False:
                 raise TicketStore.InvalidTicket('Validation failed')
 
             for k, v in kwargs.items():
                 if v is not None:
-                    dct[k] = v
+                    saved_data[k] = v
 
             # Reserialize
-            data = pickle.dumps(dct)
+            data = pickletools.optimize(
+                pickle.dumps(saved_data, protocol=-1)
+            )  # nosec: Tickets are generated by us, so we know they are safe
             if secure:
                 data = CryptoManager().aes_crypt(data, key)
             t.data = data
@@ -298,7 +307,9 @@ class TicketStore(UUIDModel):
 
     def __str__(self) -> str:
         # Tickets are generated by us, so we know they are safe
-        data = pickle.loads(self.data) if self.owner != consts.ticket.TICKET_SECURED_ONWER else '{Secure Ticket}'  # nosec
+        data = (
+            pickle.loads(self.data) if self.owner != consts.ticket.TICKET_SECURED_ONWER else '{Secure Ticket}'
+        )  # nosec
 
         return (
             f'Ticket id: {self.uuid}, Owner: {self.owner}, Stamp: {self.stamp}, '