[Python] Tip - copyreg로 pickle을 신뢰할 수 있게 하기

내장모듈 pickle은 파이썬 객체를 스트림으로 직렬화하거나 바이트를 객체로 역직렬화하는 데 사용합니다. pickle로 만든 바이트 스트림을 신뢰할 수 없는 부분과 통신하는 데 사용하면 안됩니다. pickle의 목적은 바이너리 채널을 통해 사용자가 제어하는 프로그램 간에 파이썬 객체를 넘겨주는 것입니다.

pickle 모듈의 직렬화 포맷은 설계 관점에서 안전하지 못합니다. 직렬화한 데이터는 원래의 파이썬 객체를 재구성하는 데 필요한 프로그램을 담습니다. 이는 악성 pickle 페이로드(payload)로 파이썬 프로그램에서 해당 페이로드를 역직렬화하는 부분을 망가뜨릴 수 있음을 의미합니다. 이와 달리 json 모듈은 설계 관점에서 안전합니다. 직렬화된 JSON 데이터는 객체 계층에 대한 간단한 설명을 포함합니다. JSON 데이터를 역직렬화한다고 해서 파이썬 프로그램이 추가적인 위협에 노출되지는 않습니다. JSON 같은 포맷은 서로 신뢰하지 않는 프로그램이나 사람 간에 통신 목적으로 사용해야 합니다.

예를 들어 게임에서 플레이어의 진행 상태를 파이썬 객체로 표현한다고 가정합니다. 게임 상태는 플레이어의 레벨과 남은 생명 수를 포함합니다.

class GameState: def __init__(self): self.level = 0 self.lives = 4 # 프로그램은 게임이 실행 중일 때 이 객체를 수정 state = GameState() state.level += 1 # 플레이어가 레벨을 통과함 state.lives =- 1 # 플레이어가 재도전을 해야함

사용자가 게임을 끝내면 프로그램은 게임의 상태를 파일에 저장해서 나중에 재개할 수 있게 합니다. pickle 모듈을 이용하면 이런 작업을 쉽게 할 수 있습니다. 다음은 GameState 객체를 파일에 직접 덤프하는 코드입니다.

import pickle class GameState: def __init__(self): self.level = 0 self.lives = 4 # 프로그램은 게임이 실행 중일 때 이 객체를 수정 state = GameState() state.level += 1 # 플레이어가 레벨을 통과함 state.lives -= 1 # 플레이어가 재도전을 해야함 state_path = 'game_state.bin' with open(state_path, 'wb') as f: pickle.dump(state, f) # 나중에 파일을 로드하고 직렬화한 적이 없는 것처럼 GameState 객체를 복원함 with open(state_path, 'rb') as f: state_after = pickle.load(f) print(state_after.__dict__) # 결과 # {'lives': 3, 'level': 1}

이 방법의 문제점은 시간이 지남에 따라 게임의 기능을 확장하면서 발생합니다. 플레이어가 포인트를 쌓아가게 하고 싶다고 가정합니다. 플레이어의 포인트를 추적하려고 GameState 클래스에 새로운 필드를 추가합니다.

pickle로 GameState 클래스의 새로운 버전을 직렬화하는 기능은 이전과 동일하게 동작합니다. 다음은 dumps를 사용해 객체를 문자열로 직렬화하고 loads를 사용해 객체로 되돌려서 파일을 이용한 과정을 시뮬레이트하는 코드입니다.

import pickle class GameState: def __init__(self): self.level = 0 self.lives = 4 self.points = 0 # 프로그램은 게임이 실행 중일 때 이 객체를 수정 state = GameState() serialized = pickle.dumps(state) state_after = pickle.loads(serialized) print(state_after.__dict__) # 결과 # {'lives': 4, 'level': 0, 'points': 0}

하지만 사용자가 오래 전에 저장한 GameState 객체를 이용해서 게임을 재개하면 어떻게 될까요? 다음 코드에서는 새로 정의한 GameState 클래스로 오래된 게임 파일을 언피클(unpickle)합니다.

state_path = 'game_state.bin' with open(state_path, 'rb') as f: state_after = pickle.load(f) print(state_after.__dict__) # 결과 # {'level': 1, 'lives': 3}

points 속성이 빠져서 결과가 출력됩니다. 반환된 객체가 새로 만든 GameState클래스의 인스턴스이기 때문에 당황스럽습니다.

assert isinstance(state_after, GameState)

이런 동작은 pickle 모듈이 동작할 때 생기는 부작용입니다. pickle을 사용하는 주 용도는 객체를 쉽게 직렬화하는 것입니다. pickle을 단순한 용도 이상으로 사용하자마자 모듈의 기능이 놀랍게 망가집니다. 

내장 모듈 copyreg를 이용하면 이 문제를 간단히 해결할 수 있습니다. copyreg모듈로 파이썬 객체를 직렬화할 함수를 등록하여 pickle의 동작을 제어하고 pickle을 더 신뢰할 수 있게 만들 수 있습니다.

기본 속성 값

GameState 객체가 언피클링 후에 항상 모든 속성을 담음을 보장하는 가장 간단한 방법은 기본 인수가 있는 생성자를 사용하는 것입니다. 다음은 이 방법으로 재정의한 생성자입니다.

import pickle import copyreg class GameState: def __init__(self, level=0, lives=4, points=0): self.level = level self.lives = lives self.points = points # 이 생성자를 피클용으로 사용하려고 GameState 객체를 받아 copyreg 모듈용 파라미터 튜플로 변환하는 헬퍼함수를 정의 # 반환된 튜플은 언피클링에 사용할 함수와 이 함수에 전달할 파라미터를 담음 def pickle_game_state(game_state): kwargs = game_state.__dict__ return unpickle_game_state, (kwargs,) # unpickle_game_state 헬퍼를 정의해야함. 이 함수는 직렬화된 데이터와 pickle_game_state로부터 가져온 파라미터를 받고 그에 해당하는 GameState 객체를 반환. # 이 함수는 GameState 생성자를 감싼 아주 작은 래퍼(wrapper). def unpickle_game_state(kwargs): return GameState(**kwargs) # 내장 모듈 copyreg로 GameState 객체와 직렬화 함수를 등록 copyreg.pickle(GameState, pickle_game_state) # 직렬화와 역직렬화는 이전과 동일하게 동작 state = GameState() state.points += 1000 serialized = pickle.dumps(state) state_after = pickle.loads(serialized) print(state_after.__dict__) # 결과 # {'points': 1000, 'lives': 4, 'level': 0}

등록을 마치고 나면 GameState의 정의를 변경하여 플레이어에게 사용할 마법의 개수를 부여하게 만듭니다. 이 변경은 GameState에 points 필드를 추가할 때와 비슷합니다.

class GameState: def __init__(self, level=0, lives=4, points=0, magic=5): self.level = level self.lives = lives self.points = points self.magic = magic # 추가

이번에는 이전과 달리 오래된 객체를 역직렬화해도 빠진 속성 없이 올바른 게임 데이터를 만들어냅니다. unpickle_game_state가 GameState 생성자를 직접 호출하기 때문입니다. 생성자의 키워드 인수는 파라미터가 빠지면 기본값을 가집니다. 따라서 이전 게임 상태 파일이 역직렬화될 때 새로 추가한 magic필드는 기본값을 받습니다.

state_after = pickle.loads(serialized) print(state_after.__dict__) # 결과 # {'points': 1000, 'lives': 4, 'magic': 5, 'level': 0}

요약

내장 모듈 pickle은 신뢰할 수 있는 프로그램 간에 객체를 직렬화하고 역직렬화하는 용도로만 사용할 수 있음 pickle 모듈은 간단한 사용 사례를 벗어나는 용도로 사용하면 제대로 동작하지 않을 수도 있음 빠뜨린 속성 값을 추가하거나 클래스에 버전 관리 기능을 제공하거나 안정적인 임포트 경로를 제공하려면 pickle과 함께 내장 모듈 copyreg를 사용해야 함