CLI test

Author: beskep

ECO2/cli.py

@@ -1,6 +1,6 @@
 import click
 
-from eco2 import Eco2
+from ECO2.eco2 import Eco2
 
 
 @click.group()

ECO2/eco2.py

@@ -1,6 +1,5 @@
 from itertools import cycle
 from pathlib import Path
-from typing import Union
 
 
 class Eco2:
@@ -15,7 +14,8 @@ class Eco2:
         (8, 'unknown'),
     )
     key = (172, 41, 85, 66)
-    encoding = 'UTF-8'
+    header_encoding = 'EUC-KR'
+    value_encoding = 'UTF-8'
     value_ext = '.xml'
 
     @classmethod
@@ -43,8 +43,9 @@ def _decode_header(cls, data: bytes):
         b = data
         for length, name in cls.header:
             value, b = cls._decode_chunk(b=b, length=length)
+
             try:
-                value = value.decode('euc-kr')
+                value = value.decode(cls.header_encoding)
             except ValueError:
                 pass
 
@@ -53,42 +54,39 @@ def _decode_header(cls, data: bytes):
         return header
 
     @classmethod
-    def _decrypt_eco2_data(cls, data: bytes):
-        decrypted = cls.decrypt_bytes(data)
-        hl = cls.header_length()
+    def _print_header_info(cls, header: bytes):
+        header_dict = cls._decode_header(header)
 
-        header_bytes = decrypted[:hl]
-        value_bytes = decrypted[hl:]
-        value = value_bytes.decode()
+        print('Header info:')
 
-        return header_bytes, value
+        for key, value in header_dict.items():
+            if key == 'unknown':
+                continue
+
+            print(f'    {key:10s}: {value}')
 
     @classmethod
     def _write_value(cls, path: Path, value: str):
-        path.write_text(value.replace('\r\n', '\n'), encoding=cls.encoding)
+        path.write_text(value.replace('\r\n', '\n'),
+                        encoding=cls.value_encoding)
 
     @classmethod
     def _read_value(cls, path: Path):
-        return path.read_text(encoding=cls.encoding).replace('\n', '\r\n')
+        return path.read_text(encoding=cls.value_encoding).replace('\n', '\r\n')
 
     @classmethod
-    def _print_header_info(cls, header: bytes):
-        header_dict = cls._decode_header(header)
-
-        print('Header info:')
+    def _decrypt_eco2_data(cls, data: bytes):
+        decrypted = cls.decrypt_bytes(data)
+        hl = cls.header_length()
 
-        for key, value in header_dict.items():
-            if key == 'unknown':
-                continue
+        header_bytes = decrypted[:hl]
+        value_bytes = decrypted[hl:]
+        value = value_bytes.decode()
 
-            print(f'    {key:10s}: {value}')
+        return header_bytes, value
 
     @classmethod
-    def decrypt(cls,
-                path: Union[str, Path],
-                save_dir=None,
-                header_name=None,
-                value_name=None):
+    def decrypt(cls, path, save_dir=None, header_name=None, value_name=None):
         path = Path(path)
         save_dir = path.parent if save_dir is None else Path(save_dir)
         if not save_dir.exists():
@@ -110,6 +108,13 @@ def decrypt(cls,
         value_path = save_dir.joinpath(value_name + cls.value_ext)
         cls._write_value(path=value_path, value=value)
 
+    @classmethod
+    def _encrypt(cls, header: bytes, value_path: str, save_path: Path):
+        value = cls._read_value(path=value_path)
+        bvalue = value.encode(cls.value_encoding)
+        encrypted = cls.encrypt_bytes(header + bvalue)
+        save_path.write_bytes(encrypted)
+
     @classmethod
     def encrypt(cls, header_path, value_path, save_path=None):
         if save_path is None:
@@ -119,22 +124,19 @@ def encrypt(cls, header_path, value_path, save_path=None):
         value_path = Path(value_path)
         save_path = Path(save_path)
 
-        bheader = header_path.read_bytes()
-        cls._print_header_info(bheader)
-
-        value = cls._read_value(path=value_path)
-        bvalue = value.encode()
+        header = header_path.read_bytes()
+        cls._print_header_info(header)
 
-        data = bheader + bvalue
-        encrypted = cls.encrypt_bytes(data)
-
-        save_path.write_bytes(encrypted)
+        cls._encrypt(header=header, value_path=value_path, save_path=save_path)
 
     @classmethod
     def encrypt_dir(cls, header_path, value_path, save_dir=None):
         header_path = Path(header_path)
         value_path = Path(value_path)
 
+        header = header_path.read_bytes()
+        cls._print_header_info(header)
+
         save_dir = value_path if save_dir is None else Path(save_dir)
         if not save_dir.is_dir():
             save_dir = save_dir.parent
@@ -145,6 +147,6 @@ def encrypt_dir(cls, header_path, value_path, save_dir=None):
             vps = [value_path]
 
         for vp in vps:
-            cls.encrypt(header_path=header_path,
-                        value_path=vp,
-                        save_path=save_dir.joinpath(f'{vp.stem}.eco'))
+            cls._encrypt(header=header,
+                         value_path=vp,
+                         save_path=save_dir.joinpath(f'{vp.stem}.eco'))

tests/test_eco2.py

@@ -2,17 +2,22 @@
 
 import pytest
 
+from ECO2.cli import cli
 from ECO2.eco2 import Eco2
 
+data_dir = Path(__file__).parent.joinpath('data')
+eco2_path = data_dir.joinpath('test.eco')
+header_path = data_dir.joinpath('header')
+value_path = data_dir.joinpath('value.xml')
 
-def test_eco2():
-    data_dir = Path(__file__).parent.joinpath('data')
 
-    eco2_path = data_dir.joinpath('test.eco')
-    header_path = data_dir.joinpath('header')
-    value_path = data_dir.joinpath('value.xml')
+def test_eco2():
     encrypted_path = data_dir.joinpath('encrypted.eco')
 
+    header_path.unlink(missing_ok=True)
+    value_path.unlink(missing_ok=True)
+    encrypted_path.unlink(missing_ok=True)
+
     Eco2.decrypt(path=eco2_path,
                  save_dir=data_dir,
                  header_name=header_path.stem,
@@ -34,6 +39,52 @@ def test_eco2():
 
     assert hash(eco2) == hash(encrypted)
 
+    header_path.unlink(missing_ok=True)
+    value_path.unlink(missing_ok=True)
+    encrypted_path.unlink(missing_ok=True)
+
+
+def test_cli():
+    encrypted_path = value_path.with_suffix('.eco')
+
+    header_path.unlink(missing_ok=True)
+    value_path.unlink(missing_ok=True)
+    encrypted_path.unlink(missing_ok=True)
+
+    try:
+        cli([
+            'decrypt',
+            eco2_path.as_posix(),
+            '-s',
+            eco2_path.parent.as_posix(),
+            '-h',
+            header_path.as_posix(),
+            '-v',
+            value_path.with_suffix('').as_posix(),
+        ])
+    except SystemExit:
+        pass
+
+    assert header_path.exists()
+    assert value_path.exists()
+
+    try:
+        cli([
+            'encrypt',
+            header_path.as_posix(),
+            value_path.parent.as_posix(),
+            '-s',
+            header_path.parent.as_posix(),
+        ])
+    except SystemExit:
+        pass
+
+    assert encrypted_path.exists()
+
+    header_path.unlink(missing_ok=True)
+    value_path.unlink(missing_ok=True)
+    encrypted_path.unlink(missing_ok=True)
+
 
 if __name__ == '__main__':
     pytest.main(['-v'])