# ファイルを読み込みます
def read(path):
    with open(path, mode = "rb") as f:
	    data = f.read()

    return data

# ファイルに書き込みます
def write(path, data):
    with open(path, mode = "wb") as f:
	    f.write(data)

from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
from Crypto.Util.Padding import pad
from Crypto.Util.Padding import unpad
import binary_file 

# 共通鍵の作成
def make_key(key_path):
    key_data = get_random_bytes(32) # os.urandom(32)  32 * 8 = 256ビット https://docs.python.org/ja/3/library/os.html
    binary_file.write(key_path, key_data)


# 暗号化してファイル出力します
def encrypt_file(key_path, input_path, output_path):
    key_data   = binary_file.read(key_path)   # 鍵の読み込み
    input_data = binary_file.read(input_path) # 暗号化対象データの読み込み
    encrypted_data = encrypt(key_data, input_data) # 暗号化

    # 「iv + 暗号化データ」をファイル出力します
    binary_file.write(output_path, encrypted_data) 

# 暗号化
def encrypt(key_data, input_data):
    cipher = AES.new(key_data, AES.MODE_CBC)
    encrypted_data = cipher.encrypt(pad(input_data, AES.block_size)) # パディング : pkcs7 
    return cipher.iv + encrypted_data # 「+」で bytes 同士を結合して返します


# 復号してファイル出力します
def decrypt_file(key_path, input_path, output_path):
    key_data   = binary_file.read(key_path)   # 鍵の読み込み
    input_data = binary_file.read(input_path) # 暗号化対象データの読み込み
    original_data = decrypt(key_data, input_data) # 暗号化

    # 復号結果をファイル出力します
    binary_file.write(output_path, original_data)

# 復号
def decrypt(key_data, input_data):
    # 「iv + 暗号化データ」を iv と 暗号化データに分割します
    iv             = input_data[ 0:16]
    encrypted_data = input_data[16:]

    # 復号
    cipher = AES.new(key_data, AES.MODE_CBC, iv = iv)
    return unpad(cipher.decrypt(encrypted_data), AES.block_size)

import argparse
import os
import cipher_utils

# コマンドラインの引数を取得します
def get_args():
    parser = argparse.ArgumentParser(description = '暗号化・復号')
    parser.add_argument('-v' , '--version', action='store_true', help = 'バージョン') # フラグ扱い。付けるとTrue
    parser.add_argument('-e' , '--encrypt', action='store_true', help = '暗号化')     # フラグ扱い。付けるとTrue
    parser.add_argument('-d' , '--decrypt', action='store_true', help = '復号')       # フラグ扱い。付けるとTrue
    parser.add_argument('-mk', '--makekey', help = '鍵の作成')
    parser.add_argument('-k' , '--key'    , help = '鍵')
    parser.add_argument('-i' , '--input'  , help = '入力パス')
    parser.add_argument('-o' , '--output' , help = '出力パス')
    args = parser.parse_args()

    if (args.encrypt == False and args.decrypt == False and args.makekey == None and args.version == False):
        raise Exception('--makekey、--encrypt、--decrypt のいずれかを指定してください。')

    if     (args.encrypt and args.decrypt) \
        or (args.makekey != None and args.encrypt) \
        or (args.makekey != None and args.decrypt):
        raise Exception('--makekey、--encrypt、--decrypt は、同時に指定しないでください。')

    if (args.encrypt or args.decrypt) and args.key == None:
        raise Exception('--key を指定してください。')

    if (args.encrypt or args.decrypt) and args.input == None:
        raise Exception('--input を指定してください。')

    if (args.encrypt or args.decrypt) and args.output == None:
        raise Exception('--output を指定してください。')

    return args

def main():

    # コマンドラインの引数を取得します
    try:
        args = get_args()
    except Exception as e:
        print(e) # エラーメッセージの表示
        exit(1)  # 終了します

    if args.version:
        # バージョンの表示
        print("version 1.0")
        return

    if args.makekey != None:
        # 共通鍵の作成
        if os.path.exists(args.makekey):
            print("既に同じファイル名の鍵が存在します。鍵の作成を中止します。")
            exit(1)
            
        cipher_utils.make_key(args.makekey)
        print("鍵の作成が完了しました。", args.makekey)
        return

    if args.encrypt:
        # 暗号化

        # # ファイルサイズが 1GB を超える場合は確認する
        # file_size = os.path.getsize(args.input)
        # if file_size >= 1024 * 1024 * 1024:
        #     yes_no = input("ファイルサイズが 1GB 以上あります。暗号化しますか？(y/n)")

        #     # y か yes 以外の場合は処理を中止します
        #     if not (yes_no == "y" or yes_no == "yes"):
        #         print("暗号化を中止しました。")
        #         exit(1)

        cipher_utils.encrypt_file(args.key, args.input, args.output)
        print("暗号化が完了しました。", args.output)
        return

    if args.decrypt:
        # 復号
        cipher_utils.decrypt_file(args.key, args.input, args.output)
        print("復号が完了しました。", args.output)
        return

# メイン
if __name__ == '__main__':
    main()

import unittest
import sys
import os

# VSCode からテストをする場合は、launch.json の "args" をコメントにしてから実行してください。

# src 配下のモジュールをインポートするため、src フォルダにパスを通します
py_path = os.path.abspath(__file__)
sys.path.append(os.path.dirname(os.path.dirname(py_path)) + "/src")

from aes import get_args
import cipher_utils
import binary_file

class Test_aes(unittest.TestCase):

    # 同時指定ができないコマンドラインの引数のパターン
    def test_get_args_1(self):
        py = 'aes.py'

        # 何も指定しない
        with self.assertRaises(Exception) as e:
            sys.argv = [py]
            args = get_args()

        # -e と -d を同時に指定
        with self.assertRaises(Exception) as e:
            sys.argv = [py, '-e', '-d']
            args = get_args()

        # -e と -mk を同時に指定
        with self.assertRaises(Exception) as e:
            sys.argv = [py, '-e', '-mk', 'k.dat']
            args = get_args()

        # -d と -mk を同時に指定
        with self.assertRaises(Exception) as e:
            sys.argv = [py, '-d', '-mk', 'k.dat']
            args = get_args()

    # コマンドラインの引数が足りないパターン
    def test_get_args_2(self):
        py = 'aes.py'

        # -k が未指定
        with self.assertRaises(Exception) as e:
            sys.argv = [py, '-e', '-i', 'input.txt', '-o', 'output.dat']
            args = get_args()

        # -i が未指定
        with self.assertRaises(Exception) as e:
            sys.argv = [py, '-e', '-k', 'secret.key', '-o', 'output.dat']
            args = get_args()

        # -o が未指定
        with self.assertRaises(Exception) as e:
            sys.argv = [py, '-e', '-k', 'secret.key', '-i', 'input.txt']
            args = get_args()

    # コマンドラインの引数が正常なパターン
    def test_get_args_3(self):
        py = 'aes.py'

        sys.argv = [py, '-e', '-k', 'secret.key', '-i', 'input.txt', '-o', 'output.dat']
        args = get_args()
        self.assertEqual(args.key   , 'secret.key')
        self.assertEqual(args.input , 'input.txt')
        self.assertEqual(args.output, 'output.dat')

        sys.argv = [py, '-mk', 'secret.key']
        args = get_args()
        self.assertEqual(args.makekey, 'secret.key')

    # 暗号化と復号
    def test_encrypt_decrypt(self):
        # 鍵の作成
        key_path = "test.key"
        cipher_utils.make_key(key_path)

        # 暗号化
        test_data = 'おはようございます。\nこんにちは\nこんばんは\nおやすみなさい。'
        binary_file.write("test.txt", test_data.encode('utf-8'))
        cipher_utils.encrypt_file(key_path, "test.txt", "test.dat")

        # 復号
        cipher_utils.decrypt_file(key_path, "test.dat", "test_org.txt")
        test_data_org = binary_file.read("test_org.txt").decode('utf-8')

        # テストファイルの削除
        os.remove(key_path)
        os.remove("test.txt")
        os.remove("test.dat")
        os.remove("test_org.txt")

        # 検証
        self.assertEqual(test_data, test_data_org)
        

if __name__ == '__main__':
    unittest.main()

            //鍵作成
            //"args": ["-mk", "test.key"]

            //暗号化
            //"args": ["-e", "-k", "test.key", "-i", "test.txt", "-o", "test2.txt"]

            //復号
            //"args": ["-d", "-k", "test.key", "-i", "test2.txt", "-o", "test_org.txt"]

python c:/python/aes/src/aes.py %*

set "PATH=%PATH%;C:/python"