diff --git a/README.md b/README.md index 51fdb86..81cf3f6 100644 --- a/README.md +++ b/README.md @@ -66,7 +66,7 @@ needles = [ ] haystack_file_path = './addr_list.txt' -addresses_found_list: List[str] = search_multiple_in_file( +addresses_found_list: list[str] = search_multiple_in_file( haystack_file_path=haystack_file_path, needles=needled, ) diff --git a/pyproject.toml b/pyproject.toml index 9957ee6..e1ba451 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,3 +42,7 @@ used_addr_check = "used_addr_check.cli:main_cli" [tool.pyright] typeCheckingMode = "basic" + +[tool.ruff.lint] +select = ["ALL"] +ignore = ["D", "S", "TD", "FIX", "COM812"] diff --git a/src/used_addr_check/__init__.py b/src/used_addr_check/__init__.py index 06782b2..57184d3 100644 --- a/src/used_addr_check/__init__.py +++ b/src/used_addr_check/__init__.py @@ -1,17 +1,30 @@ __VERSION__ = "0.1.6" __AUTHOR__ = "RecRanger" -from .index_create import ( # noqa F401 - load_or_generate_index, +from .cli import main_cli +from .index_create import ( generate_index, - store_index_json, load_index_json, - store_index_parquet, load_index_parquet, + load_or_generate_index, + store_index_json, + store_index_parquet, ) -from .index_search import ( # noqa F401 +from .index_search import ( search_in_file_with_index, search_multiple_in_file, # <- main library function ) -from .index_types import IndexEntry # noqa F401 -from .cli import main_cli # noqa F401 +from .index_types import IndexEntry + +__all__ = [ + "IndexEntry", + "generate_index", + "load_index_json", + "load_index_parquet", + "load_or_generate_index", + "main_cli", + "search_in_file_with_index", + "search_multiple_in_file", + "store_index_json", + "store_index_parquet", +] diff --git a/src/used_addr_check/cli.py b/src/used_addr_check/cli.py index 1ae1ebf..e757c05 100644 --- a/src/used_addr_check/cli.py +++ b/src/used_addr_check/cli.py @@ -1,16 +1,16 @@ -# from used_addr_check.download_list import download_list, BITCOIN_LIST_URL +import argparse +import sys +from pathlib import Path + from used_addr_check import __VERSION__ +from used_addr_check.defaults import DEFAULT_INDEX_CHUNK_SIZE +from used_addr_check.download_list import BITCOIN_LIST_URL, download_list from used_addr_check.index_create import load_or_generate_index from used_addr_check.index_search import search_multiple_in_file from used_addr_check.scan_file import scan_file_for_used_addresses -from used_addr_check.defaults import DEFAULT_INDEX_CHUNK_SIZE - -import argparse -import sys -from pathlib import Path -def main_cli(): +def main_cli() -> None: parser = argparse.ArgumentParser( description="CLI for file processing and searching" ) @@ -32,23 +32,21 @@ def main_cli(): subparsers = parser.add_subparsers(dest="command") # # Subparser for the 'download' command - # download_parser = subparsers.add_parser( - # "download", help="Download the file" - # ) - # download_parser.add_argument( - # "-o", - # "--output", - # dest="output_path", - # required=True, - # help="Output file path (should end in .gz)", - # ) - # download_parser.add_argument( - # "-u", - # "--url", - # dest="url", - # default=BITCOIN_LIST_URL, - # help="URL to download the file from", - # ) + download_parser = subparsers.add_parser("download", help="Download the file") + download_parser.add_argument( + "-o", + "--output", + dest="output_path", + required=True, + help="Output file path (should end in .gz)", + ) + download_parser.add_argument( + "-u", + "--url", + dest="url", + default=BITCOIN_LIST_URL, + help="URL to download the file from", + ) # Subparser for the 'version' command (subparser not really used) subparsers.add_parser( @@ -59,7 +57,10 @@ def main_cli(): # Subparser for the 'index' command index_parser = subparsers.add_parser( "index", - help="Index a haystack 'used addresses' file, save it to orig_name.txt.index.json", # noqa + help=( + "Index a haystack 'used addresses' file, " + "and save it to orig_name.txt.index.json" + ), ) index_parser.add_argument( "-f", @@ -90,7 +91,7 @@ def main_cli(): # Subparser for the 'scan_file' command scan_file_parser = subparsers.add_parser( "scan_file", - help="Scan a file for bitcoin addresses, and see which ones have been used.", # noqa + help="Scan a file for bitcoin addresses, and see which ones have been used.", ) scan_file_parser.add_argument( "-f", @@ -111,7 +112,7 @@ def main_cli(): args = parser.parse_args() if args.command == "version" or args.version: - print(f"used_addr_scan version v{__VERSION__}") + print(f"used_addr_scan version v{__VERSION__}") # noqa: T201 sys.exit(0) elif args.command == "index": load_or_generate_index( @@ -125,8 +126,8 @@ def main_cli(): args.needles, index_chunk_size=args.index_chunk_size, ) - # elif args.command == "download": - # download_list(Path(args.output_path)) + elif args.command == "download": + download_list(Path(args.output_path)) elif args.command == "scan_file": scan_file_for_used_addresses( Path(args.haystack_file_path), diff --git a/src/used_addr_check/download_list.py b/src/used_addr_check/download_list.py index 3598d16..7eed0db 100644 --- a/src/used_addr_check/download_list.py +++ b/src/used_addr_check/download_list.py @@ -1,8 +1,8 @@ from pathlib import Path import backoff -from loguru import logger import requests +from loguru import logger # TODO: This download feature isn't really production-ready. # TODO: Make the downloader as fast as wget. @@ -10,7 +10,7 @@ BITCOIN_LIST_URL = ( - "http://alladdresses.loyce.club/all_Bitcoin_addresses_ever_used_sorted.txt.gz" # noqa + "http://alladdresses.loyce.club/all_Bitcoin_addresses_ever_used_sorted.txt.gz" ) @@ -31,6 +31,7 @@ def _get_remote_file_size(url: str) -> int: ) def _download_file(url: str, dest: Path) -> Path: """Download a file from a URL to a local destination, resuming if possible. + If the file already exists and is fully downloaded, do nothing. If the file is partially downloaded, resume the download. If the file does not exist, download it from scratch. @@ -40,22 +41,16 @@ def _download_file(url: str, dest: Path) -> Path: """ total_size = _get_remote_file_size(url) - if dest.is_dir(): - destination = dest / Path(url).name - else: - destination = dest + destination = (dest / Path(url).name) if dest.is_dir() else dest - # Check how much of the file has already been downloaded - if destination.exists(): - current_size = destination.stat().st_size - else: - current_size = 0 + # Check how much of the file has already been downloaded. + current_size = destination.stat().st_size if destination.exists() else 0 # If file is already fully downloaded if current_size == total_size: logger.info("File already downloaded.") return destination - elif current_size > total_size: + if current_size > total_size: logger.warning( f"File size mismatch. Current size: {current_size:,} bytes, " f"Total size: {total_size:,} bytes" @@ -72,7 +67,7 @@ def _download_file(url: str, dest: Path) -> Path: with requests.get(url, headers=headers, stream=True) as r: r.raise_for_status() # Open the file in append mode, create if does not exist - with open(destination, "ab") as f: + with destination.open("ab") as f: for chunk in r.iter_content(chunk_size=8192): if chunk: # filter out keep-alive new chunks f.write(chunk) diff --git a/src/used_addr_check/index_create.py b/src/used_addr_check/index_create.py index 79e6249..4dfb91b 100644 --- a/src/used_addr_check/index_create.py +++ b/src/used_addr_check/index_create.py @@ -1,18 +1,17 @@ from pathlib import Path -from typing import List -from loguru import logger import orjson import polars as pl +from loguru import logger from tqdm import tqdm -from used_addr_check.index_types import IndexEntry from used_addr_check.defaults import DEFAULT_INDEX_CHUNK_SIZE +from used_addr_check.index_types import IndexEntry def generate_index( haystack_file_path: Path, index_chunk_size: int = DEFAULT_INDEX_CHUNK_SIZE -) -> List[IndexEntry]: +) -> list[IndexEntry]: """ Generates an index for a large sorted text file, storing every `index_chunk_size` lines. @@ -25,7 +24,7 @@ def generate_index( - List[IndexEntry]: A list of tuples containing line text, byte offset, and line number. """ - index: List[IndexEntry] = [] + index: list[IndexEntry] = [] haystack_file_size = haystack_file_path.stat().st_size with ( tqdm( @@ -35,7 +34,7 @@ def generate_index( total=haystack_file_size, desc="Indexing haystack file", ) as progress_bar, - open(haystack_file_path, "rb") as file, + haystack_file_path.open("rb") as file, ): offset = 0 for line_number, line in enumerate(file): @@ -56,7 +55,7 @@ def generate_index( return index -def store_index_json(index: List[IndexEntry], index_json_file_path: Path) -> None: +def store_index_json(index: list[IndexEntry], index_json_file_path: Path) -> None: """ Stores the index in a file for later use. @@ -65,11 +64,11 @@ def store_index_json(index: List[IndexEntry], index_json_file_path: Path) -> Non - index_file_path (Path): The path to store the index. """ - with open(index_json_file_path, "wb") as file: + with index_json_file_path.open("wb") as file: file.write(orjson.dumps(index)) -def load_index_json(index_json_file_path: Path) -> List[IndexEntry]: +def load_index_json(index_json_file_path: Path) -> list[IndexEntry]: """ Loads an index from a JSON file. @@ -79,12 +78,12 @@ def load_index_json(index_json_file_path: Path) -> List[IndexEntry]: Returns: - List[IndexEntry]: The loaded index. """ - with open(index_json_file_path, "rb") as file: - raw_read: List[dict] = orjson.loads(file.read()) + with index_json_file_path.open("rb") as file: + raw_read: list[dict] = orjson.loads(file.read()) return [IndexEntry(**val) for val in raw_read] -def store_index_parquet(index: List[IndexEntry], index_parquet_file_path: Path) -> None: +def store_index_parquet(index: list[IndexEntry], index_parquet_file_path: Path) -> None: """ Stores the index in a parquet file for later use. @@ -97,7 +96,7 @@ def store_index_parquet(index: List[IndexEntry], index_parquet_file_path: Path) df.write_parquet(index_parquet_file_path) -def load_index_parquet(index_parquet_file_path: Path) -> List[IndexEntry]: +def load_index_parquet(index_parquet_file_path: Path) -> list[IndexEntry]: """ Loads an index from a parquet file. @@ -114,8 +113,9 @@ def load_index_parquet(index_parquet_file_path: Path) -> List[IndexEntry]: def load_or_generate_index( haystack_file_path: Path, index_chunk_size: int = DEFAULT_INDEX_CHUNK_SIZE, + *, force_recreate: bool = False, -) -> List[IndexEntry]: +) -> list[IndexEntry]: """Attempts to load an index from a file, or generates one if it doesn't, or if `force_recreate` is enabled. @@ -153,4 +153,5 @@ def load_or_generate_index( logger.info(f"Index loaded with {len(index):,} entries") return index - raise Exception("This should be unreachable.") + msg = "This should be unreachable." + raise RuntimeError(msg) diff --git a/src/used_addr_check/index_search.py b/src/used_addr_check/index_search.py index 1bd87a3..9577601 100644 --- a/src/used_addr_check/index_search.py +++ b/src/used_addr_check/index_search.py @@ -1,17 +1,16 @@ import bisect -from pathlib import Path -from typing import List import json +from pathlib import Path from loguru import logger from tqdm import tqdm +from used_addr_check.defaults import DEFAULT_INDEX_CHUNK_SIZE from used_addr_check.index_create import load_or_generate_index from used_addr_check.index_types import IndexEntry -from used_addr_check.defaults import DEFAULT_INDEX_CHUNK_SIZE -def _binary_search_index(index: List[IndexEntry], needle: str) -> int: +def _binary_search_index(index: list[IndexEntry], needle: str) -> int: """ Performs a binary search on the index to find the closest position for the needle string. @@ -34,18 +33,20 @@ def _binary_search_index(index: List[IndexEntry], needle: str) -> int: if index[idx].line_value > needle: # common occurrence return idx - 1 - elif index[idx].line_value == needle: + if index[idx].line_value == needle: # rare occurrence - exact match on one of the index keys (e.g., 0) return idx - elif index[idx].line_value < needle: + if index[idx].line_value < needle: # not sure if this can ever happen - raise Exception("bisect.bisect_left failed to find the correct position") - else: - raise Exception("Unexpected condition") + msg = "bisect.bisect_left failed to find the correct position" + raise RuntimeError(msg) + + msg = "Unexpected condition" + raise RuntimeError(msg) def search_in_file_with_index( - haystack_file_path: Path, needle: str, index: List[IndexEntry] + haystack_file_path: Path, needle: str, index: list[IndexEntry] ) -> bool: """ Searches for a needle string in the file using a pre-built index to narrow @@ -63,12 +64,13 @@ def search_in_file_with_index( assert isinstance(index, list) position = _binary_search_index(index, needle) - # logger.debug(f"Search {needle}: Binary search position: {position}") - # logger.debug(f"Search {needle}: {index[position]}") + + if 0: + logger.debug(f"Search {needle}: Binary search position: {position}") + logger.debug(f"Search {needle}: {index[position]}") + if position == len(index): - # logger.debug( - # "Binary search position equals index length, returning False" - # ) + # Binary search position equals index length, returning False. return False # Find the bounds to search within the file @@ -77,7 +79,7 @@ def search_in_file_with_index( if position + 1 < len(index): end_offset = index[position + 1].byte_offset - with open(haystack_file_path, "r", encoding="ascii") as file: + with haystack_file_path.open(encoding="ascii") as file: file.seek(start_offset) while True: if end_offset and file.tell() >= end_offset: @@ -92,9 +94,9 @@ def search_in_file_with_index( def search_multiple_in_file( haystack_file_path: Path | str, - needles: List[str] | str, + needles: list[str] | str, index_chunk_size: int = DEFAULT_INDEX_CHUNK_SIZE, -) -> List[str]: +) -> list[str]: """ Searches for multiple needle strings in the file by pre-building an index and then searching within the file. @@ -113,11 +115,13 @@ def search_multiple_in_file( index = load_or_generate_index(haystack_file_path, index_chunk_size) - # do the search - found_needles = [] - for needle in tqdm(needles, desc="Searching needles", unit="needle"): - if search_in_file_with_index(haystack_file_path, needle, index=index): - found_needles.append(needle) + # Do the search. + found_needles = [ + needle + for needle in tqdm(needles, desc="Searching needles", unit="needle") + if search_in_file_with_index(haystack_file_path, needle, index=index) + ] + logger.info(f"Found {len(found_needles):,}/{len(needles):,} needles in the file") logger.info(f"Needles found: {json.dumps(sorted(found_needles))}") return found_needles diff --git a/src/used_addr_check/scan_file.py b/src/used_addr_check/scan_file.py index b6de052..e3fbfc0 100644 --- a/src/used_addr_check/scan_file.py +++ b/src/used_addr_check/scan_file.py @@ -1,20 +1,19 @@ -from pathlib import Path -from typing import List, Literal import re +from collections.abc import Sequence +from pathlib import Path +from typing import Literal from loguru import logger -from ripgrepy import Ripgrepy, RipGrepNotFound +from ripgrepy import RipGrepNotFound, Ripgrepy -from used_addr_check.index_search import search_multiple_in_file from used_addr_check.defaults import DEFAULT_INDEX_CHUNK_SIZE - -# BITCOIN_ADDR_REGEX = r"[13][a-km-zA-HJ-NP-Z1-9]{25,34}" +from used_addr_check.index_search import search_multiple_in_file # Source: https://ihateregex.io/expr/bitcoin-address/ BITCOIN_ADDR_REGEX = r"\b((bc1|[13])[a-zA-HJ-NP-Z0-9]{25,39})\b" -def _extract_addresses_from_file_python_re(text_file_path: Path) -> List[str]: +def _extract_addresses_from_file_python_re(text_file_path: Path) -> list[str]: """ Extracts bitcoin addresses from a file using Python regex. @@ -33,7 +32,7 @@ def _extract_addresses_from_file_python_re(text_file_path: Path) -> List[str]: return [result[0] for result in results] -def _extract_addresses_from_file_ripgrep(text_file_path: Path) -> List[str]: +def _extract_addresses_from_file_ripgrep(text_file_path: Path) -> list[str]: """ Extracts bitcoin addresses from a file using ripgrep. @@ -72,11 +71,11 @@ def _extract_addresses_from_file_ripgrep(text_file_path: Path) -> List[str]: def extract_addresses_from_file( text_file_path: Path, - enabled_searchers: List[Literal["ripgrep", "python_re"]] = [ + enabled_searchers: Sequence[Literal["ripgrep", "python_re"]] = ( "ripgrep", "python_re", - ], -) -> List[str]: + ), +) -> list[str]: """ Extracts bitcoin addresses from a file, using either ripgrep or Python re. @@ -107,16 +106,18 @@ def extract_addresses_from_file( return _extract_addresses_from_file_python_re(text_file_path) else: - raise ValueError(f"Invalid searcher provided: {searcher}") + msg = f"Invalid searcher provided: {searcher}" + raise ValueError(msg) - raise Exception("This should never be reached. Address extraction has failed.") + msg = "This should never be reached. Address extraction has failed." + raise RuntimeError(msg) def scan_file_for_used_addresses( haystack_file_path: Path, needle_file_path: Path, index_chunk_size: int = DEFAULT_INDEX_CHUNK_SIZE, -): +) -> None: """ Scans a file for bitcoin addresses, and see which one have been used. diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_scan_file.py b/tests/test_scan_file.py index 06c2efc..8b8b017 100644 --- a/tests/test_scan_file.py +++ b/tests/test_scan_file.py @@ -1,4 +1,6 @@ +from collections.abc import Callable from pathlib import Path + import pytest from used_addr_check.scan_file import ( @@ -29,7 +31,7 @@ # Combine tests for multiple functions and files @pytest.mark.parametrize( - "func, input_file, expected", + ("func", "input_file", "expected"), [ ( _extract_addresses_from_file_python_re, @@ -63,6 +65,8 @@ ), ], ) -def test_extract_addresses_from_file(func, input_file, expected): +def test_extract_addresses_from_file( + func: Callable[[Path], list[str]], input_file: Path, expected: list[str] +) -> None: actual = func(input_file) assert actual == expected