RSA를 이용한, 메시지 암호화 & 복호화

# =============================================================================
# 공개키 메시지를 이용한 메시지/파일 암복호화
# =============================================================================
import base64

from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_OAEP

import tools

# =============================================================================
#  key 생성 및 로딩
# =============================================================================
def make_key(pr_path, pu_path, size_bit=1024):
    pr_key = RSA.generate(size_bit)
    pu_key = pr_key.publickey()
    
    fw = open(pr_path, 'wb')
    fw.write(pr_key.exportKey('PEM'))
    fw.close()
    
    fw = open(pu_path,'wb')
    fw.write(pu_key.exportKey('PEM'))
    fw.close()

#make_key('pr1.key','pu1.key')

def get_key(key_path):
    fr = open(key_path, 'rb')
    key = RSA.importKey(fr.read())
    return key

# =============================================================================
#  메시지 암복호화
# =============================================================================
def encrypt_msg(msg, key):
    ch = PKCS1_OAEP.new(key) #암호를 푼다.
    msg_en = ch.encrypt(msg)
    
    return msg_en

def decrypt_msg(msg, key):
    ch = PKCS1_OAEP.new(key)
    msg_de = ch.decrypt(msg)
    
    return msg_de

def test1():
    msg = 'i love you you'
    
    pu_key = get_key('pu1.key')
    msg_en = encrypt_msg(msg.encode(), pu_key)
    pr_key = get_key('pr1.key')
    msg_de = decrypt_msg(msg_en, pr_key)

    print("암호화 : ", msg_en)
    print("길이 : ", len(msg_en))
    print("복호화 : ", msg_de)

#test1()

# =============================================================================
# 파일 암복호화 
# =============================================================================

def encrypt_file(file_name, key):
    save_name = file_name + '.enc'
    
    fr = open(file_name,'rb')
    fw = open(save_name, 'wb')
    
    BSIZE = 86# 최대 86까지 가능함!
   # pu_key = get_key('pu1_key')    
    
    data = fr.read(BSIZE)
    while data :
        data_enc = encrypt_msg(data, key)
        fw.write(data_enc)
        data = fr.read()
        
    return True
    
def decrypt_file(file_name, key):
    save_file = file_name + '.txt'
    
    fr = open(file_name, 'rb')
    fw = open(save_file, 'wb')
    
    BSIZE = 128 # 입력값과 상관없이 128이므로 128단위로 읽어야 함. 즉, key 길이에 따라 바뀜
    data = fr.read(BSIZE)
    
    while data:
        data_dec = decrypt_msg(data, key)
        fw.write(data_dec)
        data = fr.read()
        
    return True


def test2():
    file_name = 'my1.py'
    pu_key = get_key('pu1.key')
    encrypt_file(file_name, pu_key)
    pr_key = get_key('pr1.key')
    file_name = 'my1.py.enc'
    decrypt_file(file_name, pr_key)
    
# =============================================================================
#  해시검증
# =============================================================================
    hash1 = tools.get_hash_file('my1.py') 
    hash2 = tools.get_hash_file('my1.py.enc.txt')
    
    if(hash1 == hash2):
        print("SAME")
    else:
        print("NOT SAME!")
    
test2()