Skip to content

encryption

checkpoint.encryption

체크포인트 암호화 관련 함수들을 정의합니다.

torch.save로 저장 가능한 파이썬 오브젝트를 암호화하거나 복호화합니다.

  • encrypt: torch.save로 저장 가능한 파이썬 오브젝트를 암호화합니다.
  • decrypt: torch.load로 로드 가능한 바이트를 복호화합니다.
  • encrypt_dict: dict 자체 혹은 요소들을 암호화합니다.
  • decrypt_dict: dict를 복호화합니다.
  • encrypt_save_dict: dict를 암호화 후 torch.save 로 저장합니다.
  • decrypt_load_dict: torch.load로 파일 로드 후 복호화해 dict로 로드합니다.

BUFFER_SIZE module-attribute 

BUFFER_SIZE = 64 * 1024

handle_unpickling_error 

handle_unpickling_error(function)
Source code in SaigeToolkit/checkpoint/metadata/metadata.py 
16
17
18
19
20
21
22
23
24
def handle_unpickling_error(function):
    @wraps(function)
    def wrapper(*args, **kwargs):
        try:
            return function(*args, **kwargs)
        except pickle.UnpicklingError:
            raise InvalidModelFileError

    return wrapper

encrypt 

encrypt(obj: Any, password: str, buffer_size: int = BUFFER_SIZE, **torch_save_args) -> bytes

torch.save로 picklable한 obj를 Aes 암호화 후 바이트로 리턴합니다.

Parameters:

  • obj (Any) –

    description

  • password (str) –

    description

  • buffer_size (int, default: BUFFER_SIZE ) –

    description. Defaults to BUFFER_SIZE.

Returns:

  • bytes ( bytes ) –

    description

Source code in SaigeToolkit/checkpoint/encryption.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
def encrypt(obj: Any, password: str, buffer_size: int = BUFFER_SIZE, **torch_save_args) -> bytes:
    """`torch.save`로 picklable한 obj를 Aes 암호화 후 바이트로 리턴합니다.

    Args:
        obj (Any): _description_
        password (str): _description_
        buffer_size (int, optional): _description_. Defaults to BUFFER_SIZE.

    Returns:
        bytes: _description_
    """
    buffer = io.BytesIO()
    torch.save(obj, buffer, **torch_save_args)
    buffer.seek(0)
    buffer_encrypt = io.BytesIO()
    pyAesCrypt.encryptStream(buffer, buffer_encrypt, password, buffer_size)
    return buffer_encrypt.getvalue()

decrypt 

decrypt(encrpyted: bytes, password: str, buffer_size: int = BUFFER_SIZE, **torch_load_args) -> Any

torch.load로 로드 가능한 바이트를 Aes 복호화 후 파이썬 오브젝트로 리턴합니다.

Parameters:

  • encrpyted (bytes) –

    description

  • password (str) –

    description

  • buffer_size (int, default: BUFFER_SIZE ) –

    description. Defaults to BUFFER_SIZE.

Returns:

  • Any ( Any ) –

    description

Source code in SaigeToolkit/checkpoint/encryption.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def decrypt(encrpyted: bytes, password: str, buffer_size: int = BUFFER_SIZE, **torch_load_args) -> Any:
    """`torch.load`로 로드 가능한 바이트를 Aes 복호화 후 파이썬 오브젝트로 리턴합니다.

    Args:
        encrpyted (bytes): _description_
        password (str): _description_
        buffer_size (int, optional): _description_. Defaults to BUFFER_SIZE.

    Returns:
        Any: _description_
    """
    decrypt_buffer = io.BytesIO()
    pyAesCrypt.decryptStream(
        io.BytesIO(encrpyted), decrypt_buffer, password, buffer_size, len(encrpyted)
    )
    decrypt_buffer.seek(0)
    return torch.load(decrypt_buffer, **torch_load_args)

encrypt_dict 

encrypt_dict(data: Mapping, password: Optional[str] = None, keys_encrypt: Optional[List[str]] = None, **torch_save_args)

dict 자체 혹은 요소들을 암호화합니다.

Parameters:

  • data (Mapping) –

    description

  • password (Optional[str], default: None ) –

    description. Defaults to None.

  • keys_encrypt (Optional[List[str]], default: None ) –

    None이 아닐 경우 keys_encrypt에 있는 키의 값들만 암호화합니다. Defaults to None.

Source code in SaigeToolkit/checkpoint/encryption.py 
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
def encrypt_dict(
    data: Mapping,
    password: Optional[str] = None,
    keys_encrypt: Optional[List[str]] = None,
    **torch_save_args,
):
    """dict 자체 혹은 요소들을 암호화합니다.

    Args:
        data (Mapping): _description_
        password (Optional[str], optional): _description_. Defaults to None.
        keys_encrypt (Optional[List[str]], optional): None이 아닐 경우 keys_encrypt에 있는 키의 값들만 암호화합니다. Defaults to None.
    """
    data_ = data.copy()
    if password is not None:
        if keys_encrypt is not None:
            for key in keys_encrypt:
                if key not in data_:
                    continue
                data_[key] = encrypt(data_[key], password, **torch_save_args)
        else:
            data_ = encrypt(data_, password, **torch_save_args)
    return data_

decrypt_dict 

decrypt_dict(data, password: Optional[str] = None, keys_encrypt: Optional[List[str]] = None, **torch_load_args) -> Mapping

dict를 복호화합니다.

Parameters:

  • data (Any) –

    description

  • password (Optional[str], default: None ) –

    description. Defaults to None.

  • keys_encrypt (Optional[List[str]], default: None ) –

    None이 아닐 경우 keys_encrypt에 있는 키의 값들만 복호화합니다. Defaults to None.

Raises:

  • FileNotFoundError –

    description

Returns:

  • Mapping ( Mapping ) –

    description

Source code in SaigeToolkit/checkpoint/encryption.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
def decrypt_dict(
    data,
    password: Optional[str] = None,
    keys_encrypt: Optional[List[str]] = None,
    **torch_load_args,
) -> Mapping:
    """dict를 복호화합니다.

    Args:
        data (Any): _description_
        password (Optional[str], optional): _description_. Defaults to None.
        keys_encrypt (Optional[List[str]], optional): None이 아닐 경우 keys_encrypt에 있는 키의 값들만 복호화합니다. Defaults to None.

    Raises:
        FileNotFoundError: _description_

    Returns:
        Mapping: _description_
    """
    if password is not None:
        if keys_encrypt is not None:
            for key in keys_encrypt:
                if key not in data:
                    continue
                data[key] = decrypt(data[key], password, **torch_load_args)
        else:
            data = decrypt(data, password, **torch_load_args)
    return data

encrypt_save_dict 

encrypt_save_dict(data: Mapping, path: str, password: Optional[str] = None, keys_encrypt: Optional[List[str]] = None, **torch_save_args) -> None

dict를 암호화 후 torch.save 로 저장합니다.

Parameters:

  • data (Mapping) –

    description

  • path (str) –

    description

  • password (Optional[str], default: None ) –

    description. Defaults to None.

  • keys_encrypt (Optional[List[str]], default: None ) –

    None이 아닐 경우 keys_encrypt에 있는 키의 값들만 암호화합니다. Defaults to None.

Source code in SaigeToolkit/checkpoint/encryption.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
def encrypt_save_dict(
    data: Mapping,
    path: str,
    password: Optional[str] = None,
    keys_encrypt: Optional[List[str]] = None,
    **torch_save_args,
) -> None:
    """dict를 암호화 후 `torch.save` 로 저장합니다.

    Args:
        data (Mapping): _description_
        path (str): _description_
        password (Optional[str], optional): _description_. Defaults to None.
        keys_encrypt (Optional[List[str]], optional): None이 아닐 경우 keys_encrypt에 있는 키의 값들만 암호화합니다. Defaults to None.
    """
    data_ = encrypt_dict(data, password=password, keys_encrypt=keys_encrypt, **torch_save_args)
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "wb") as f:
        torch.save(data_, f, **torch_save_args)

decrypt_load_dict 

decrypt_load_dict(path: str, password: Optional[str] = None, keys_encrypt: Optional[List[str]] = None, **torch_load_args) -> Mapping

torch.load로 파일 로드 후 복호화해 dict로 로드합니다.

Parameters:

  • path (str) –

    description

  • password (Optional[str], default: None ) –

    description. Defaults to None.

  • keys_encrypt (Optional[List[str]], default: None ) –

    None이 아닐 경우 keys_encrypt에 있는 키의 값들만 복호화합니다. Defaults to None.

Raises:

  • FileNotFoundError –

    description

Returns:

  • Mapping ( Mapping ) –

    description

Source code in SaigeToolkit/checkpoint/encryption.py 
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
@handle_unpickling_error
def decrypt_load_dict(
    path: str,
    password: Optional[str] = None,
    keys_encrypt: Optional[List[str]] = None,
    **torch_load_args,
) -> Mapping:
    """`torch.load`로 파일 로드 후 복호화해 dict로 로드합니다.

    Args:
        path (str): _description_
        password (Optional[str], optional): _description_. Defaults to None.
        keys_encrypt (Optional[List[str]], optional): None이 아닐 경우 keys_encrypt에 있는 키의 값들만 복호화합니다. Defaults to None.

    Raises:
        FileNotFoundError: _description_

    Returns:
        Mapping: _description_
    """
    if not os.path.isfile(path):
        raise error.FileNotFoundError
    data = torch.load(path, **torch_load_args)
    data = decrypt_dict(data, password=password, keys_encrypt=keys_encrypt, **torch_load_args)
    return data