diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f83cd51..a497dce 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -12,55 +12,56 @@ jobs: strategy: matrix: os: [ubuntu-latest, windows-latest, macos-latest] - python-version: ['3.13'] + python-version: ['3.9', '3.10', '3.11', '3.12', '3.13', '3.14'] steps: - uses: actions/checkout@v4 - + - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - + allow-prereleases: true + - name: Install dependencies run: | python -m pip install --upgrade pip pip install pytest pytest-cov - + - name: Run tests run: | python -m pytest test_ksuid.py -v --cov=. --cov-report=xml - + - name: Run benchmarks run: | python benchmark.py - + - name: Test CLI run: | - python -m ksuid.cli generate --count 5 - python -m ksuid.cli benchmark --count 1000 + python cli.py generate --count 5 + python cli.py benchmark --count 1000 lint: runs-on: ubuntu-latest - + steps: - uses: actions/checkout@v4 - + - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: '3.13' - + - name: Install dependencies run: | python -m pip install --upgrade pip pip install black flake8 mypy - + - name: Run black run: black --check . - + - name: Run flake8 run: flake8 . --max-line-length=88 --extend-ignore=E203,W503 - + - name: Run mypy run: mypy . --ignore-missing-imports diff --git a/__init__.py b/__init__.py index 9c4f5b6..d45fd70 100644 --- a/__init__.py +++ b/__init__.py @@ -20,14 +20,15 @@ """ import os +import secrets import time from datetime import datetime, timezone from typing import Union, Optional __version__ = "1.0.0" -__all__ = ["KSUID", "generate", "from_string", "from_bytes"] +__all__ = ["KSUID", "generate", "generate_token", "from_string", "from_bytes"] -# KSUID epoch (January 1, 2014 UTC) +# KSUID epoch (May 13, 2014 16:53:20 UTC) EPOCH = 1400000000 # KSUID components @@ -43,14 +44,16 @@ class KSUID: """ K-Sortable Unique Identifier - + A KSUID is a 20-byte identifier consisting of: - 4-byte timestamp (seconds since KSUID epoch) - 16-byte random payload - + KSUIDs are naturally sortable by creation time and collision-resistant. """ - + + __slots__ = ('_timestamp', '_payload', '_bytes') + def __init__(self, timestamp: Optional[int] = None, payload: Optional[bytes] = None): """ Create a new KSUID. @@ -147,7 +150,7 @@ def __repr__(self) -> str: def __eq__(self, other) -> bool: if not isinstance(other, KSUID): - return False + return NotImplemented return self._bytes == other._bytes def __lt__(self, other) -> bool: @@ -178,13 +181,10 @@ def _base62_encode(data: bytes) -> str: """Encode bytes to base62 string.""" if not data: return "" - + # Convert bytes to integer num = int.from_bytes(data, 'big') - - if num == 0: - return BASE62_ALPHABET[0] - + result = [] while num > 0: num, remainder = divmod(num, BASE62_BASE) @@ -196,17 +196,27 @@ def _base62_encode(data: bytes) -> str: return encoded.zfill(27) +_BASE62_LOOKUP = {c: i for i, c in enumerate(BASE62_ALPHABET)} + +# Maximum integer value that fits in TOTAL_LENGTH bytes +_MAX_ENCODED = (1 << (TOTAL_LENGTH * 8)) - 1 + + def _base62_decode(s: str) -> bytes: """Decode base62 string to bytes.""" if not s: return b"" - + num = 0 for char in s: - if char not in BASE62_ALPHABET: + val = _BASE62_LOOKUP.get(char) + if val is None: raise ValueError(f"Invalid base62 character: {char}") - num = num * BASE62_BASE + BASE62_ALPHABET.index(char) - + num = num * BASE62_BASE + val + + if num > _MAX_ENCODED: + raise ValueError("Base62 value exceeds maximum for KSUID") + # Convert to bytes (20 bytes for KSUID) return num.to_bytes(TOTAL_LENGTH, 'big') @@ -217,6 +227,20 @@ def generate() -> KSUID: return KSUID() +def generate_token() -> str: + """Generate a cryptographically secure opaque token as a base62 string. + + Unlike KSUIDs, tokens use 20 bytes (160 bits) of pure random data from + ``secrets.token_bytes`` with no embedded timestamp. This makes them + suitable for API keys, session secrets, and other security-sensitive + values where the creation time should not be leaked. + + Returns: + A 27-character base62 string with 160 bits of entropy. + """ + return _base62_encode(secrets.token_bytes(TOTAL_LENGTH)) + + def from_string(ksuid_str: str) -> KSUID: """Create a KSUID from its string representation.""" return KSUID.from_string(ksuid_str) diff --git a/cli.py b/cli.py index 552af70..094d23c 100644 --- a/cli.py +++ b/cli.py @@ -14,6 +14,24 @@ from datetime import datetime +MAX_COUNT = 1_000_000 + + +def _validate_count(value): + """Validate --count is a positive integer within bounds.""" + try: + ivalue = int(value) + except ValueError: + raise argparse.ArgumentTypeError(f"invalid integer value: {value!r}") + if ivalue < 1: + raise argparse.ArgumentTypeError("count must be at least 1") + if ivalue > MAX_COUNT: + raise argparse.ArgumentTypeError( + f"count must be at most {MAX_COUNT:,}" + ) + return ivalue + + def cmd_generate(args): """Generate one or more KSUIDs.""" for _ in range(args.count): @@ -127,10 +145,10 @@ def main(): # Generate command gen_parser = subparsers.add_parser('generate', help='Generate KSUIDs') gen_parser.add_argument( - '-c', '--count', - type=int, - default=1, - help='Number of KSUIDs to generate (default: 1)' + '-c', '--count', + type=_validate_count, + default=1, + help='Number of KSUIDs to generate (default: 1, max: 1,000,000)' ) gen_parser.add_argument( '-v', '--verbose', @@ -158,10 +176,10 @@ def main(): # Benchmark command bench_parser = subparsers.add_parser('benchmark', help='Run benchmark') bench_parser.add_argument( - '-c', '--count', - type=int, - default=10000, - help='Number of KSUIDs to generate (default: 10000)' + '-c', '--count', + type=_validate_count, + default=10000, + help='Number of KSUIDs to generate (default: 10000, max: 1,000,000)' ) bench_parser.set_defaults(func=cmd_benchmark) diff --git a/prefixed_examples.py b/prefixed_examples.py index eefaa8c..071e3b7 100644 --- a/prefixed_examples.py +++ b/prefixed_examples.py @@ -10,7 +10,7 @@ import os sys.path.insert(0, os.path.dirname(__file__)) -from __init__ import KSUID, generate, from_string +from __init__ import KSUID, generate, generate_token, from_string from typing import Dict, Optional, Tuple import re @@ -87,9 +87,9 @@ def create(cls, prefix: str) -> str: if not prefix: raise ValueError("Prefix cannot be empty") - # Validate prefix format (alphanumeric and underscores only) - if not re.match(r'^[a-zA-Z][a-zA-Z0-9_]*$', prefix): - raise ValueError("Prefix must start with a letter and contain only alphanumeric characters and underscores") + # Validate prefix format (alphanumeric only, no underscores since _ is the delimiter) + if not re.match(r'^[a-zA-Z][a-zA-Z0-9]*$', prefix): + raise ValueError("Prefix must start with a letter and contain only alphanumeric characters") return f"{prefix}_{generate()}" @@ -194,12 +194,20 @@ def create_order_id() -> str: return PrefixedKSUID.create('ord') def create_api_key() -> str: - """Create an API key: ak_...""" - return PrefixedKSUID.create('ak') + """Create a secure API key: ak_... + + Uses 160 bits of cryptographically secure random data (no timestamp) + via ``generate_token()``, making it safe for use as a secret key. + """ + return f"ak_{generate_token()}" def create_session_id() -> str: - """Create a session ID: sess_...""" - return PrefixedKSUID.create('sess') + """Create a secure session token: sess_... + + Uses 160 bits of cryptographically secure random data (no timestamp) + via ``generate_token()``, making it safe for use as a bearer token. + """ + return f"sess_{generate_token()}" def demo_basic_usage(): diff --git a/test_ksuid.py b/test_ksuid.py index 6f519c7..157e088 100644 --- a/test_ksuid.py +++ b/test_ksuid.py @@ -9,7 +9,7 @@ import pytest from datetime import datetime, timezone -from __init__ import KSUID, generate, from_string, from_bytes, EPOCH +from __init__ import KSUID, generate, generate_token, from_string, from_bytes, EPOCH class TestKSUID: @@ -90,6 +90,11 @@ def test_from_string_invalid_characters(self): """Test that invalid base62 characters raise error.""" with pytest.raises(ValueError, match="Invalid base62 character"): KSUID.from_string("!" * 27) # Invalid character + + def test_from_string_overflow(self): + """Test that a base62 string exceeding 20-byte max raises ValueError.""" + with pytest.raises(ValueError, match="Base62 value exceeds maximum"): + KSUID.from_string("z" * 27) # Exceeds 2^160 - 1 def test_from_bytes(self): """Test creating KSUID from bytes.""" @@ -134,13 +139,20 @@ def test_equality(self): payload = b'\x01' * 16 ksuid1 = KSUID(timestamp=timestamp, payload=payload) ksuid2 = KSUID(timestamp=timestamp, payload=payload) - + assert ksuid1 == ksuid2 assert hash(ksuid1) == hash(ksuid2) - + # Different payload should not be equal ksuid3 = KSUID(timestamp=timestamp, payload=b'\x02' * 16) assert ksuid1 != ksuid3 + + def test_equality_with_non_ksuid(self): + """Test that __eq__ returns NotImplemented for non-KSUID types.""" + ksuid = KSUID() + assert ksuid.__eq__("not a ksuid") is NotImplemented + assert ksuid.__eq__(42) is NotImplemented + assert ksuid.__eq__(None) is NotImplemented def test_string_representation(self): """Test string and repr methods.""" @@ -241,6 +253,115 @@ def test_base62_encoding_properties(self): assert len(ksuid_str) == 27 +class TestGenerateToken: + """Test cases for generate_token() secure token generation.""" + + def test_token_is_27_chars(self): + """Token string is exactly 27 base62 characters.""" + token = generate_token() + assert len(token) == 27 + + def test_token_is_base62(self): + """Token contains only valid base62 characters.""" + valid = set("0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz") + token = generate_token() + assert all(c in valid for c in token) + + def test_tokens_are_unique(self): + """Multiple tokens must all be distinct.""" + tokens = {generate_token() for _ in range(200)} + assert len(tokens) == 200 + + def test_token_differs_from_ksuid_structure(self): + """Token's first 4 bytes should not match a KSUID timestamp. + + A real KSUID encodes (time.time() - EPOCH) in its first 4 bytes. + A pure-random token should almost never land in the same narrow + range. We generate 10 tokens and verify none of them decode to + a timestamp within 1 year of now (probability ≈ (2/2^32)^10). + """ + from __init__ import _base62_decode + now = int(time.time()) + one_year = 365 * 86400 + for _ in range(10): + raw = _base62_decode(generate_token()) + ts_value = int.from_bytes(raw[:4], 'big') + EPOCH + if abs(ts_value - now) > one_year: + return # At least one clearly non-timestamp token — pass + pytest.fail("All 10 tokens decoded to timestamps near 'now' — extremely unlikely") + + +class TestSlots: + """Verify KSUID uses __slots__ for memory efficiency.""" + + def test_no_instance_dict(self): + """KSUID instances should not have a __dict__.""" + ksuid = KSUID() + assert not hasattr(ksuid, '__dict__') + + def test_cannot_set_arbitrary_attribute(self): + """Setting an undefined attribute should raise AttributeError.""" + ksuid = KSUID() + with pytest.raises(AttributeError): + ksuid.foo = "bar" + + +class TestThreadSafety: + """Verify KSUID generation is safe under concurrent threads.""" + + def test_concurrent_generate(self): + """KSUIDs generated across threads must all be unique.""" + from concurrent.futures import ThreadPoolExecutor + count_per_thread = 500 + num_threads = 4 + + def gen_batch(_): + return [generate() for _ in range(count_per_thread)] + + with ThreadPoolExecutor(max_workers=num_threads) as pool: + batches = list(pool.map(gen_batch, range(num_threads))) + + all_ksuids = [k for batch in batches for k in batch] + assert len(all_ksuids) == count_per_thread * num_threads + assert len(set(all_ksuids)) == len(all_ksuids) + + def test_concurrent_generate_token(self): + """Tokens generated across threads must all be unique.""" + from concurrent.futures import ThreadPoolExecutor + count_per_thread = 500 + num_threads = 4 + + def gen_batch(_): + return [generate_token() for _ in range(count_per_thread)] + + with ThreadPoolExecutor(max_workers=num_threads) as pool: + batches = list(pool.map(gen_batch, range(num_threads))) + + all_tokens = [t for batch in batches for t in batch] + assert len(all_tokens) == count_per_thread * num_threads + assert len(set(all_tokens)) == len(all_tokens) + + +class TestEdgeCases: + """Edge cases for encoding / decoding.""" + + def test_zero_value_round_trip(self): + """All-zero KSUID must encode to 27 chars and round-trip correctly.""" + ksuid = KSUID(timestamp=EPOCH, payload=b'\x00' * 16) + s = str(ksuid) + assert len(s) == 27 + assert s == '0' * 27 + assert KSUID.from_string(s) == ksuid + + def test_max_timestamp_round_trip(self): + """KSUID at the maximum timestamp must round-trip correctly.""" + max_ts = EPOCH + 2**32 - 1 + ksuid = KSUID(timestamp=max_ts, payload=b'\xff' * 16) + s = str(ksuid) + assert len(s) == 27 + assert KSUID.from_string(s) == ksuid + + if __name__ == "__main__": # Run tests if script is executed directly import sys