Satisfy linter

Author: Sneagan

generate-address-list.py

@@ -1,64 +1,109 @@
 #!/usr/bin/env python3
 
-import xml.etree.ElementTree as ET
 import argparse
-import pathlib
-import json
 import hashlib
+import json
+import pathlib
+import xml.etree.ElementTree as ET
+
 from ofac_scraper import OfacWebsiteScraper
 
 FEATURE_TYPE_TEXT = "Digital Currency Address - "
-NAMESPACE = {'sdn': 'https://sanctionslistservice.ofac.treas.gov/api/PublicationPreview/exports/ADVANCED_XML'}
+NAMESPACE = {
+    'sdn': 'https://sanctionslistservice.ofac.treas.gov/api/PublicationPreview/exports/'
+    'ADVANCED_XML'
+}
 
 # List of implemented output formats
 OUTPUT_FORMATS = ["TXT", "JSON"]
 
 SDN_ADVANCED_FILE_PATH = "sdn_advanced.xml"
 
+
 def parse_arguments():
     parser = argparse.ArgumentParser(
-        description='Tool to extract sanctioned digital currency addresses from the OFAC special designated nationals XML file (sdn_advanced.xml)')
-    parser.add_argument('assets', nargs='*',
-                        default=[], help='the asset for which the sanctioned addresses should be extracted (default: XBT (Bitcoin))')
-    parser.add_argument('-sdn', '--special-designated-nationals-list', dest='sdn', type=argparse.FileType('rb'),
-                        help='the path to the sdn_advanced.xml file (can be downloaded from https://www.treasury.gov/ofac/downloads/sanctions/1.0/sdn_advanced.xml)', default=SDN_ADVANCED_FILE_PATH)
-    parser.add_argument('-f', '--output-format',  dest='format', nargs='*', choices=OUTPUT_FORMATS,
-                        default=OUTPUT_FORMATS[0], help='the output file format of the address list (default: TXT)')
-    parser.add_argument('-path', '--output-path', dest='outpath',  type=pathlib.Path, default=pathlib.Path(
-        "./"), help='the path where the lists should be written to (default: current working directory ("./")')
+        description='Tool to extract sanctioned digital currency addresses from the '
+        'OFAC special designated nationals XML file (sdn_advanced.xml)')
+    parser.add_argument(
+        'assets',
+        nargs='*',
+        default=[],
+        help='the asset for which the sanctioned addresses should be extracted '
+        '(default: XBT (Bitcoin))'
+    )
+    parser.add_argument(
+        '-sdn',
+        '--special-designated-nationals-list',
+        dest='sdn',
+        type=argparse.FileType('rb'),
+        help='the path to the sdn_advanced.xml file (can be downloaded from '
+        'https://www.treasury.gov/ofac/downloads/sanctions/1.0/sdn_advanced.xml)',
+        default=SDN_ADVANCED_FILE_PATH
+    )
+    parser.add_argument(
+        '-f',
+        '--output-format',
+        dest='format',
+        nargs='*',
+        choices=OUTPUT_FORMATS,
+        default=OUTPUT_FORMATS[0],
+        help='the output file format of the address list (default: TXT)'
+    )
+    parser.add_argument(
+        '-path',
+        '--output-path',
+        dest='outpath',
+        type=pathlib.Path,
+        default=pathlib.Path("./"),
+        help='the path where the lists should be written to (default: current working '
+        'directory ("./")'
+    )
     return parser.parse_args()
 
+
 def feature_type_text(asset):
     """returns text we expect in a <FeatureType></FeatureType> tag for a given asset"""
     return "Digital Currency Address - " + asset
 
+
 def get_possible_assets(root):
     """
     Returns a list of possible digital currency assets from the parsed XML.
     """
     assets = []
-    feature_types = root.findall('sdn:ReferenceValueSets/sdn:FeatureTypeValues/sdn:FeatureType', NAMESPACE)
+    feature_types = root.findall(
+        'sdn:ReferenceValueSets/sdn:FeatureTypeValues/sdn:FeatureType',
+        NAMESPACE
+    )
     for feature_type in feature_types:
         if feature_type.text.startswith('Digital Currency Address - '):
             asset = feature_type.text.replace('Digital Currency Address - ', '')
             assets.append(asset)
     return assets
 
+
 def get_address_id(root, asset):
     """returns the feature id of the given asset"""
     feature_type = root.find(
-        "sdn:ReferenceValueSets/sdn:FeatureTypeValues/*[.='{}']".format(feature_type_text(asset)), NAMESPACE)
-    if feature_type == None:
-        raise LookupError("No FeatureType with the name {} found".format(
-            feature_type_text(asset)))
+        f"sdn:ReferenceValueSets/sdn:FeatureTypeValues"
+        f"/*[.='{feature_type_text(asset)}']",
+        NAMESPACE
+    )
+    if feature_type is None:
+        raise LookupError(
+            f"No FeatureType with the name {feature_type_text(asset)} found"
+        )
     address_id = feature_type.attrib["ID"]
     return address_id
 
 
 def get_sanctioned_addresses(root, address_id):
     """returns a list of sanctioned addresses for the given address_id"""
     addresses = list()
-    for feature in root.findall("sdn:DistinctParties//*[@FeatureTypeID='{}']".format(address_id), NAMESPACE):
+    for feature in root.findall(
+        f"sdn:DistinctParties//*[@FeatureTypeID='{address_id}']",
+        NAMESPACE
+    ):
         for version_detail in feature.findall(".//sdn:VersionDetail", NAMESPACE):
             addresses.append(version_detail.text)
     return addresses
@@ -72,14 +117,15 @@ def write_addresses(addresses, asset, output_formats, outpath):
 
 
 def write_addresses_txt(addresses, asset, outpath):
-    with open("{}/sanctioned_addresses_{}.txt".format(outpath, asset), 'w') as out:
+    with open(f"{outpath}/sanctioned_addresses_{asset}.txt", 'w') as out:
         for address in addresses:
-            out.write(address+"\n")
+            out.write(address + "\n")
 
 
 def write_addresses_json(addresses, asset, outpath):
-    with open("{}/sanctioned_addresses_{}.json".format(outpath, asset), 'w') as out:
-        out.write(json.dumps(addresses, indent=2)+"\n")
+    with open(f"{outpath}/sanctioned_addresses_{asset}.json", 'w') as out:
+        out.write(json.dumps(addresses, indent=2) + "\n")
+
 
 def compute_sha256(file_path):
     sha256_hash = hashlib.sha256()
@@ -88,10 +134,12 @@ def compute_sha256(file_path):
             sha256_hash.update(chunk)
     return sha256_hash.hexdigest()
 
+
 def write_checksum_file(sha256, checksum_file_path):
     with open(checksum_file_path, "w") as checksum_file:
         checksum_file.write(f"SHA256({SDN_ADVANCED_FILE_PATH}) = {sha256}\n")
 
+
 def main():
     args = parse_arguments()
 
@@ -110,7 +158,7 @@ def main():
     root = tree.getroot()
 
     assets = list()
-    if type(args.assets) == str:
+    if isinstance(args.format, str):
         assets.append(args.assets)
     else:
         assets = args.assets
@@ -119,7 +167,7 @@ def main():
         assets = get_possible_assets(root)
 
     output_formats = list()
-    if type(args.format) == str:
+    if isinstance(args.format, str):
         output_formats.append(args.format)
     else:
         output_formats = args.format
@@ -150,5 +198,6 @@ def main():
 
     write_checksum_file(sha256_checksum_from_site, "data/sdn_advanced_checksum.txt")
 
+
 if __name__ == "__main__":
     main()

ofac_scraper.py

@@ -1,14 +1,13 @@
-import logging
 import time
+
 from selenium import webdriver
 from selenium.common.exceptions import TimeoutException
-from selenium.webdriver.chrome.service import Service
 from selenium.webdriver.chrome.options import Options
+from selenium.webdriver.chrome.service import Service
 from selenium.webdriver.common.action_chains import ActionChains
 from selenium.webdriver.common.by import By
-from selenium.webdriver.support.ui import WebDriverWait
 from selenium.webdriver.support import expected_conditions as EC
-from webdriver_manager.chrome import ChromeDriverManager
+from selenium.webdriver.support.ui import WebDriverWait
 
 
 class OfacWebsiteScraper:
@@ -41,7 +40,8 @@ def get_sha256_checksum(self):
 
         for attempt in range(MAX_RETRIES):
             print(
-                f"Attempting to get SHA-256 checksum (attempt {attempt + 1}/{MAX_RETRIES})..."
+                f"Attempting to get SHA-256 checksum (attempt"
+                f" {attempt + 1}/{MAX_RETRIES})..."
             )
             try:
                 self.open_website("https://sanctionslist.ofac.treas.gov/Home/SdnList")
@@ -53,7 +53,10 @@ def get_sha256_checksum(self):
 
                 # Scroll to (waiting for animation) and Click the 'File
                 # Signatures' button with the known ID
-                header_element = self.driver.find_element(By.ID, "accordion__heading-:r1:")
+                header_element = self.driver.find_element(
+                    By.ID,
+                    "accordion__heading-:r1:"
+                )
                 ActionChains(self.driver).move_to_element(header_element).perform()
                 time.sleep(1)
                 header_element.click()
@@ -78,7 +81,8 @@ def get_sha256_checksum(self):
 
             except TimeoutException:
                 print(
-                    f"Timeout occurred on attempt {attempt + 1}/{MAX_RETRIES}. Retrying in {RETRY_DELAY} seconds..."
+                    f"Timeout occurred on attempt {attempt + 1}/{MAX_RETRIES}. "
+                    f"Retrying in {RETRY_DELAY} seconds..."
                 )
                 time.sleep(RETRY_DELAY)
                 if attempt == MAX_RETRIES - 1:

update_s3_objects.py

@@ -1,15 +1,14 @@
 #!/usr/bin/env python3
 
 import argparse
+import base64
 import concurrent.futures
 import glob
+import json
 import logging
 import os
-import json
 import threading
 from io import BytesIO
-import base64
-from typing import List
 
 import boto3
 from botocore.exceptions import ClientError
@@ -34,7 +33,8 @@
 def parse_arguments():
     """Parse command line arguments."""
     parser = argparse.ArgumentParser(
-        description='Read sanctioned addresses from TXT files and create/remove empty S3 objects to match the SDN'
+        description='Read sanctioned addresses from TXT files and create/remove empty '
+        'S3 objects to match the SDN'
     )
     parser.add_argument(
         '-d', '--directory',
@@ -123,7 +123,10 @@ def read_sanctioned_addresses(directory):
         try:
             with open(file_path) as f:
                 addresses = [line.strip() for line in f if line.strip()]
-                logger.info(f"Read {len(addresses)} addresses from {os.path.basename(file_path)}")
+                logger.info(
+                    f"Read {len(addresses)} addresses from "
+                    f"{os.path.basename(file_path)}"
+                )
                 unique_addresses.update(addresses)
         except Exception as e:
             logger.error(f"Error reading {file_path}: {e}")
@@ -145,7 +148,9 @@ def create_s3_object(address, bucket, prefix, dry_run, s3_client):
     object_key = f"{prefix}{encode(address)}"
 
     if dry_run:
-        logger.info(f"DRY RUN: Would create S3 object s3://{bucket}/{object_key} ({address})")
+        logger.info(
+            f"DRY RUN: Would create S3 object s3://{bucket}/{object_key} "
+            f"({address})")
         return True, None
 
     try:
@@ -167,7 +172,9 @@ def delete_s3_object(address, bucket, prefix, dry_run, s3_client):
     """Delete an S3 object for the given address."""
     object_key = f"{prefix}{encode(address)}"
     if dry_run:
-        logger.info(f"DRY RUN: Would delete S3 object s3://{bucket}/{object_key} ({address})")
+        logger.info(
+            f"DRY RUN: Would delete S3 object s3://{bucket}/{object_key} ({address})"
+        )
         return True, None
     try:
         s3_client.delete_object(
@@ -177,7 +184,8 @@ def delete_s3_object(address, bucket, prefix, dry_run, s3_client):
         # Check if the object was actually deleted
         try:
             s3_client.head_object(Bucket=bucket, Key=object_key)
-            return False, f"Failed to delete S3 object for {address}: Object still exists"
+            return False, f"Failed to delete S3 object for {address}: "
+            "Object still exists"
         except ClientError as e:
             if e.response['Error']['Code'] == '404':
                 # 404. File is gone
@@ -208,7 +216,8 @@ def process_action_chunk(action_chunk, bucket, prefix, dry_run, s3_client):
                     results['created'] += 1
                 else:
                     results['errors'] += 1
-                    if error: logger.error(error)
+                    if error:
+                        logger.error(error)
             case 'remove':
                 success, error = delete_s3_object(
                     action['address'], bucket, prefix, dry_run, s3_client
@@ -239,10 +248,14 @@ def reconcile_s3(
     action_list = list(actions)
     total_actions = len(action_list)
 
-    logger.info(f"Starting to take {total_actions} actions on S3 objects using {workers} worker threads")
+    logger.info(f"Starting to take {total_actions} actions on S3 objects using "
+                f"{workers} worker threads")
 
     # Create chunks of actions to process
-    action_chunks = [action_list[i:i + chunk_size] for i in range(0, total_actions, chunk_size)]
+    action_chunks = [
+        action_list[i:i + chunk_size]
+        for i in range(0, len(action_list), chunk_size)
+    ]
 
     created_count = 0
     removed_count = 0
@@ -267,9 +280,11 @@ def reconcile_s3(
                 error_count += results['errors']
 
                 logger.info(
-                    f"Completed chunk {chunk_index+1}/{len(action_chunks)}, "
-                    f"total progress: {created_count + removed_count + error_count}/{total_actions} "
-                    f"({(created_count + removed_count + error_count) / total_actions * 100:.1f}%)"
+                    f"Completed chunk {chunk_index + 1}/{len(action_chunks)}, "
+                    f"total progress: {created_count + removed_count + error_count}/"
+                    f"{total_actions} ({(
+                        created_count + removed_count + error_count
+                    ) / total_actions * 100:.1f}%)"
                 )
 
             except Exception as e:

update_s3_objects_test.py

@@ -2,22 +2,22 @@
 Test S3 interaction
 """
 import base64
-import unittest
-from unittest.mock import Mock, patch
-from typing import List
 import tempfile
+import unittest
 from pathlib import Path
+from unittest.mock import Mock, patch
+
 from botocore.exceptions import ClientError
 
 from update_s3_objects import (
+    create_s3_object,
     decode,
+    delete_s3_object,
     encode,
+    format_result_message,
     generate_actions,
-    read_sanctioned_addresses,
     process_action_chunk,
-    create_s3_object,
-    delete_s3_object,
-    format_result_message
+    read_sanctioned_addresses,
 )
 
 
@@ -89,7 +89,7 @@ def test_roundtrip(self) -> None:
 
     def test_padding_handling(self) -> None:
         """Test that the functions correctly handle padding."""
-        test_cases: List[str] = []
+        test_cases: list[str] = []
         test_cases.append("abc")  # 3 bytes = 4 base64 chars, no padding
         test_cases.append("abcd")  # 4 bytes = 6 base64 chars, 2 padding chars
         test_cases.append("abcde")  # 5 bytes = 8 base64 chars, 1 padding char