{security enhancement} Check ckpt file before loading

[cibernicola]

I found this code to check a ckpt file in order to see if it is safe or not to work with it.
I think it would be interesting to work on this (if it has not already been done) since the proliferation of this type of files and the more than possible bad intention of some actors, could present problems in the not too distant future.
For my part I am testing outside WebUI, although I still need to know exactly all the cases in which a ckpt file can be dangerous for the system that executes it.

I think it could be a python file external to WebUI, which would load and execute the different actions from the startup scripts, this way there would be a total separation and potential prevention of "attacks".
For now I have made a couple of modifications to this code to make it work automatically, loading all the ckpt files and check them in a simple loop.

PD: It would also be interesting to develop something similar for embeedings, HyperNetworks and aestethics.

Shall we give stable diffusion webui a layer of security?

import io
import zipfile
import argparse
import pickle
import collections
import torch
import numpy
import _codecs

def encode(*args):
    out = _codecs.encode(*args)
    print(f'encode({args}) = {out}')
    return out

class RestrictedUnpickler(pickle.Unpickler):
    def persistent_load(self, saved_id):
        assert saved_id[0] == 'storage'
        return torch.ByteStorage()

    def find_class(self, module, name):
        print(f'find class {module} {name}')
        if module == 'collections' and name == 'OrderedDict':
            return getattr(collections, name)
        if module == 'torch._utils' and name == '_rebuild_tensor_v2':
            return torch._utils._rebuild_tensor_v2
        if module == 'torch' and name in ['FloatStorage', 'HalfStorage']:
            return torch.FloatStorage
        if module == 'numpy.core.multiarray' and name == 'scalar':
            return numpy.core.multiarray.scalar
        if module == 'numpy' and name == 'dtype':
            return numpy.dtype
        if module == '_codecs' and name == 'encode':
            return encode
        # Forbid everything else.
        raise pickle.UnpicklingError("global '%s/%s' is forbidden" %
                                     (module, name))

def restricted_loads(s):
    """Helper function analogous to pickle.loads()."""
    return RestrictedUnpickler(io.BytesIO(s)).load()

# To test that it catches this RCE:
# restricted_loads(b"cos\nsystem\n(S'echo hello world'\ntR.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('file', type=str)
    args = parser.parse_args()

    # unzip model.ckpt archive/data.pkl
    with zipfile.ZipFile(args.file, 'r') as ckf:
        with ckf.open('archive/data.pkl', 'r') as f:
            st = f.read()
    d = restricted_loads(st)
    print(dir(d))
    print(d.keys())
    print(d['callbacks'])

src: https://rentry.org/safeunpickle2

[cibernicola]

Well, seems already implemented:
https://github.com/AUTOMATIC1111/stable-diffusion-webui/search?q=pickle&type=commits