"""Fix _sanitize_filename in storage_service.py to use chr(0) instead of literal null bytes."""
import os
import re
import ast

path = "/var/www/kelthos/platform/nomazra/storage_service.py"
with open(path, "r") as f:
    content = f.read()

lines = content.split('\n')
func_start = None
func_end = None
for i, line in enumerate(lines):
    if 'def _sanitize_filename' in line:
        func_start = i
    elif func_start is not None and (line.startswith('def ') or line.startswith('\ndef ')) and i > func_start + 2:
        func_end = i
        break

if func_start is None:
    print("ERROR: _sanitize_filename not found")
    exit(1)

if func_end is None:
    # Find next function after _sanitize_filename
    for i in range(func_start + 3, len(lines)):
        stripped = lines[i].rstrip()
        if stripped.startswith('def ') and not stripped.startswith('def _sanitize_filename'):
            func_end = i
            break
    if func_end is None:
        func_end = len(lines)

print(f"Found _sanitize_filename at lines {func_start+1}-{func_end}")

new_func = '''def _sanitize_filename(name):
    """
    Strip path traversal attacks from filenames.
    Blocks: ../ | ..\\\\ | null bytes | encoded traversal | absolute paths | control chars.
    Returns only the basename with safe characters.
    """
    if not name or not isinstance(name, str):
        raise ValueError('Invalid filename')
    # Reject null bytes (common injection vector)
    if chr(0) in name:
        raise ValueError('Invalid filename: null bytes')
    # Decode percent-encoded traversal attempts (%2e%2e%2f, %2e%2e%5c)
    decoded = name
    for _ in range(3):  # Multi-layer decode
        prev = decoded
        decoded = decoded.replace('%2e', '.').replace('%2E', '.')
        decoded = decoded.replace('%2f', '/').replace('%2F', '/')
        decoded = decoded.replace('%5c', '\\\\').replace('%5C', '\\\\')
        decoded = decoded.replace('%00', '')
        if decoded == prev:
            break
    # Take only the final filename component (strip any directory path)
    basename = decoded.replace('\\\\', '/').split('/')[-1]
    # Strip leading dots (prevent .htaccess, ..hidden, etc.)
    basename = basename.lstrip('.')
    # Remove control characters (0x00-0x1f, 0x7f)
    basename = re.sub(r'[\\x00-\\x1f\\x7f]', '', basename)
    # Only allow safe characters: alphanumeric, dash, underscore, dot, space
    basename = re.sub(r'[^a-zA-Z0-9._\\- ]', '_', basename)
    # Collapse multiple dots/underscores
    basename = re.sub(r'\\.{2,}', '.', basename)
    basename = re.sub(r'_{2,}', '_', basename)
    basename = basename.strip('. _')
    if not basename:
        raise ValueError('Invalid filename: empty after sanitization')
    # Max length
    if len(basename) > 255:
        ext = os.path.splitext(basename)[1][:20]
        basename = basename[:235] + ext
    return basename

'''

lines[func_start:func_end] = new_func.split('\n')
content = '\n'.join(lines)

with open(path, "w") as f:
    f.write(content)

# Verify no null bytes
data = open(path, "rb").read()
null_count = data.count(b"\x00")
print(f"Null bytes remaining: {null_count}")

# Syntax check
try:
    ast.parse(content)
    print("Syntax OK")
except SyntaxError as e:
    print(f"Syntax ERROR: {e}")
