style: fix lint errors

Author: djosix

src/padding_oracle/__init__.py

@@ -27,3 +27,16 @@
     to_bytes, to_str,
 )
 from .legacy import padding_oracle
+
+__all__ = [
+    'solve',
+    'convert_to_bytes',
+    'remove_padding',
+    'padding_oracle',
+    'urlencode',
+    'urldecode',
+    'base64_encode',
+    'base64_decode',
+    'to_bytes',
+    'to_str',
+]

src/padding_oracle/legacy.py

@@ -22,61 +22,75 @@
 
 import logging
 import traceback
-from typing import Optional, Union
+from typing import List, Union
 
 from .encoding import to_bytes
-from .solve import Fail, OracleFunc, ResultType, solve, remove_padding, convert_to_bytes
+from .solve import (
+    solve, Fail, OracleFunc, ResultType,
+    convert_to_bytes, remove_padding)
 
 __all__ = [
     'padding_oracle',
 ]
 
+
 def padding_oracle(ciphertext: Union[bytes, str],
                    block_size: int,
                    oracle: OracleFunc,
                    num_threads: int = 1,
                    log_level: int = logging.INFO,
                    null_byte: bytes = b' ',
                    return_raw: bool = False,
-                   ) -> bytes:
+                   ) -> Union[bytes, List[int]]:
     '''
-    Run padding oracle attack to decrypt ciphertext given a function to check wether the
-    ciphertext can be decrypted successfully.
+    Run padding oracle attack to decrypt ciphertext given a function to check
+    wether the ciphertext can be decrypted successfully.
 
     Args:
         cipher         (bytes|str) the cipher you want to decrypt
-        block_size     (int)       block size (the cipher length should be multiple of this)
+        block_size     (int)       block size (the cipher length should be
+                                   multiple of this)
         oracle         (function)  a function: oracle(cipher: bytes) -> bool
-        num_threads    (int)       how many oracle functions will be run in parallel (default: 1)
+        num_threads    (int)       how many oracle functions will be run in
+                                   parallel (default: 1)
         log_level      (int)       log level (default: logging.INFO)
-        null_byte      (bytes|str) the default byte when plaintext are not set (default: None)
-        return_raw     (bool)      do not convert plaintext into bytes and unpad (default: False)
+        null_byte      (bytes|str) the default byte when plaintext are not
+                                   set (default: None)
+        return_raw     (bool)      do not convert plaintext into bytes and
+                                   unpad (default: False)
 
     Returns:
         plaintext (bytes|List[int]) the decrypted plaintext
     '''
 
     # Check args
-    assert callable(oracle), 'the oracle function should be callable'
-    assert isinstance(ciphertext, (bytes, str)), 'cipher should have type bytes'
-    assert isinstance(block_size, int), 'block_size should have type int'
-    assert len(ciphertext) % block_size == 0, 'cipher length should be multiple of block size'
-    assert 1 <= num_threads <= 1000, 'num_threads should be in [1, 1000]'
-    assert isinstance(null_byte, (bytes, str)), 'expect null with type bytes or str'
-    assert len(null_byte) == 1, 'null byte should have length of 1'
+    if not callable(oracle):
+        raise TypeError('the oracle function should be callable')
+    if not isinstance(ciphertext, (bytes, str)):
+        raise TypeError('cipher should have type bytes')
+    if not isinstance(block_size, int):
+        raise TypeError('block_size should have type int')
+    if not len(ciphertext) % block_size == 0:
+        raise ValueError('cipher length should be multiple of block size')
+    if not 1 <= num_threads <= 1000:
+        raise ValueError('num_threads should be in [1, 1000]')
+    if not isinstance(null_byte, (bytes, str)):
+        raise TypeError('expect null with type bytes or str')
+    if not len(null_byte) == 1:
+        raise ValueError('null byte should have length of 1')
 
     logger = get_logger()
     logger.setLevel(log_level)
-    
+
     ciphertext = to_bytes(ciphertext)
     null_byte = to_bytes(null_byte)
-    
+
     # Wrapper to handle exceptions from the oracle function
     def wrapped_oracle(ciphertext: bytes):
         try:
             return oracle(ciphertext)
         except Exception as e:
-            logger.error('error calling oracle with {!r}'.format(ciphertext))
+            logger.error(f'error in oracle with {ciphertext!r}, {e}')
             logger.debug('error details: {}'.format(traceback.format_exc()))
         return False
 
@@ -93,13 +107,14 @@ def plaintext_callback(plaintext: bytes):
 
     plaintext = solve(ciphertext, block_size, wrapped_oracle, num_threads,
                       result_callback, plaintext_callback)
-    
+
     if not return_raw:
         plaintext = convert_to_bytes(plaintext, null_byte)
         plaintext = remove_padding(plaintext)
-        
+
     return plaintext
 
+
 def get_logger():
     logger = logging.getLogger('padding_oracle')
     formatter = logging.Formatter('[%(asctime)s][%(levelname)s] %(message)s')

src/padding_oracle/solve.py

@@ -22,7 +22,7 @@
 
 import asyncio
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any, Callable, NamedTuple, Set, Union, List
+from typing import Callable, NamedTuple, Set, Union, List
 
 from .encoding import to_bytes
 
@@ -33,16 +33,19 @@
     'remove_padding',
 ]
 
+
 class Pass(NamedTuple):
     block_index: int
     index: int
     byte: int
 
+
 class Fail(NamedTuple):
     block_index: int
     message: str
     is_critical: bool = False
 
+
 class Done(NamedTuple):
     block_index: int
     C0: List[int]
@@ -54,26 +57,28 @@ class Done(NamedTuple):
 OracleFunc = Callable[[bytes], bool]
 ResultCallback = Callable[[ResultType], bool]
 PlainTextCallback = Callable[[List[int]], bool]
- 
+
 
 class Context(NamedTuple):
     block_size: int
     oracle: OracleFunc
-    
-    executor: ThreadPoolExecutor 
+
+    executor: ThreadPoolExecutor
     loop: asyncio.AbstractEventLoop
-    
+
     tasks: Set[asyncio.Task[ResultType]]
 
     latest_plaintext: List[int]
     plaintext: List[int]
-    
+
     result_callback: ResultCallback
     plaintext_callback: PlainTextCallback
 
+
 def dummy_callback(*a, **ka):
     pass
 
+
 def solve(ciphertext: bytes,
           block_size: int,
           oracle: OracleFunc,
@@ -87,6 +92,7 @@ def solve(ciphertext: bytes,
                          result_callback, plaintext_callback)
     return loop.run_until_complete(future)
 
+
 async def solve_async(ciphertext: bytes,
                       block_size: int,
                       oracle: OracleFunc,
@@ -96,43 +102,47 @@ async def solve_async(ciphertext: bytes,
                       ) -> List[int]:
 
     ciphertext = list(ciphertext)
-    assert len(ciphertext) % block_size == 0, \
-        'ciphertext length must be a multiple of block_size'
-    assert len(ciphertext) // block_size > 1, \
-        'cannot solve with only one block'
+
+    if not len(ciphertext) % block_size == 0:
+        raise ValueError('ciphertext length must be a multiple of block_size')
+    if not len(ciphertext) // block_size > 1:
+        raise ValueError('cannot solve with only one block')
 
     ctx = create_solve_context(ciphertext, block_size, oracle, parallel,
                                result_callback, plaintext_callback)
 
     while True:
-        done_tasks, _ = await asyncio.wait(ctx.tasks, return_when=asyncio.FIRST_COMPLETED)
-        
+        done_tasks, _ = await asyncio.wait(ctx.tasks,
+                                           return_when=asyncio.FIRST_COMPLETED)
+
         for task in done_tasks:
             result = await task
-            
+
             ctx.result_callback(result)
             ctx.tasks.remove(task)
-            
+
             if isinstance(result, Pass):
-                update_latest_plaintext(ctx, result.block_index, result.index, result.byte)
+                update_latest_plaintext(
+                    ctx, result.block_index, result.index, result.byte)
             if isinstance(result, Done):
                 update_plaintext(ctx, result.block_index, result.C0, result.X1)
-        
+
         if len(ctx.tasks) == 0:
             break
-    
+
     # Check if any block failed
     error_block_indices = set()
-    
+
     for i, byte in enumerate(ctx.plaintext):
         if byte is None:
             error_block_indices.add(i // block_size + 1)
-            
+
     for idx in error_block_indices:
         result_callback(Fail(idx, f'cannot decrypt cipher block {idx}', True))
-        
+
     return ctx.plaintext
 
+
 def create_solve_context(ciphertext, block_size, oracle, parallel,
                          result_callback, plaintext_callback) -> Context:
     tasks = set()
@@ -143,30 +153,27 @@ def create_solve_context(ciphertext, block_size, oracle, parallel,
 
     plaintext = [None] * (len(cipher_blocks) - 1) * block_size
     latest_plaintext = plaintext.copy()
-    
+
     executor = ThreadPoolExecutor(parallel)
     loop = asyncio.get_running_loop()
     ctx = Context(block_size, oracle, executor, loop, tasks,
                   latest_plaintext, plaintext,
                   result_callback, plaintext_callback)
-    
+
     for i in range(1, len(cipher_blocks)):
         run_block_task(ctx, i, cipher_blocks[i-1], cipher_blocks[i], [])
 
     return ctx
 
+
 def run_block_task(ctx: Context, block_index, C0, C1, X1):
     future = solve_block(ctx, block_index, C0, C1, X1)
     task = ctx.loop.create_task(future)
     ctx.tasks.add(task)
 
-async def solve_block(
-    ctx: Context,
-    block_index: int,
-    C0: List[int],
-    C1: List[int],
-    X1: List[int] = [],
-) -> ResultType:
+
+async def solve_block(ctx: Context, block_index: int, C0: List[int],
+                      C1: List[int], X1: List[int] = []) -> ResultType:
     # X1 = decrypt(C1)
     # P1 = xor(C0, X1)
 
@@ -195,46 +202,56 @@ async def solve_block(
     for byte in hits:
         X1_test = [byte ^ padding, *X1]
         run_block_task(ctx, block_index, C0, C1, X1_test)
-        
+
     return Pass(block_index, index, byte ^ padding ^ C0[index])
 
-async def get_oracle_hits(ctx: Context, C0: List[int], C1: List[int], index: int):
-    
+
+async def get_oracle_hits(ctx: Context, C0: List[int], C1: List[int],
+                          index: int):
+
     C0 = C0.copy()
     futures = {}
-    
+
     for byte in range(256):
         C0[index] = byte
         ciphertext = bytes(C0 + C1)
         futures[byte] = ctx.loop.run_in_executor(
             ctx.executor, ctx.oracle, ciphertext)
-        
+
     hits = []
-    
+
     for byte, future in futures.items():
         is_valid = await future
         if is_valid:
             hits.append(byte)
-            
+
     return hits
 
-def update_latest_plaintext(ctx: Context, block_index: int, index: int, byte: int):
+
+def update_latest_plaintext(ctx: Context, block_index: int, index: int,
+                            byte: int):
+
     i = (block_index - 1) * ctx.block_size + index
     ctx.latest_plaintext[i] = byte
     ctx.plaintext_callback(ctx.latest_plaintext)
 
-def update_plaintext(ctx: Context, block_index: int, C0: List[int], X1: List[int]):
+
+def update_plaintext(ctx: Context, block_index: int, C0: List[int],
+                     X1: List[int]):
+
     assert len(C0) == len(X1) == ctx.block_size
     block = compute_plaintext(C0, X1)
-    
+
     i = (block_index - 1) * ctx.block_size
     ctx.latest_plaintext[i:i+ctx.block_size] = block
     ctx.plaintext[i:i+ctx.block_size] = block
     ctx.plaintext_callback(ctx.plaintext)
 
+
 def compute_plaintext(C0: List[int], X1: List[int]):
     return [c ^ x for c, x in zip(C0, X1)]
 
+
 def convert_to_bytes(byte_list: List[int], replacement=b' '):
     '''
     Convert a list of int into bytes, replace invalid byte with replacement.
@@ -249,6 +266,7 @@ def convert_to_bytes(byte_list: List[int], replacement=b' '):
         byte_list[i] = byte
     return bytes(byte_list)
 
+
 def remove_padding(data: Union[str, bytes, List[int]]) -> bytes:
     '''
     Remove PKCS#7 padding bytes.

tests/__init__.py

@@ -1,8 +1,11 @@
-import sys, os
+import sys
+import os
+
 
 def get_src_dir():
     current_dir = os.path.dirname(__file__)
     src_dir = os.path.join(current_dir, '..', 'src')
     return src_dir
 
+
 sys.path.insert(0, get_src_dir())