Merge pull request #1 from tonyzorin/claude/security-audit-Vwtzv

Claude/security audit vwtzv

Author: tonyzorin

.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

__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)

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)
 

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():