#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
CUPS 打印机驱动管理器
轻量级Web服务，用于上传和安装打印机驱动
"""

import os
import sys
import subprocess
import tempfile
import shutil
import tarfile
import zipfile
import gzip
import ipaddress
import uuid
from pathlib import Path
from functools import wraps

from flask import Flask, request, render_template_string, redirect, url_for, flash, jsonify, Response
from werkzeug.utils import secure_filename

app = Flask(__name__)
app.secret_key = os.urandom(24)

def env_int(name, default):
    try:
        return max(1, int(os.environ.get(name, default)))
    except ValueError:
        return default

# 配置
UPLOAD_FOLDER = '/tmp/cups-drivers'
MAX_UPLOAD_MB = env_int('DRIVER_MANAGER_MAX_UPLOAD_MB', 512)
MAX_CONTENT_LENGTH = MAX_UPLOAD_MB * 1024 * 1024
COMMAND_TIMEOUT = env_int('DRIVER_MANAGER_COMMAND_TIMEOUT', 900)
ALLOWED_EXTENSIONS = {'deb', 'ppd', 'gz', 'tar', 'tgz', 'zip', 'rpm', 'sh', 'run'}

# 管理员凭据（可通过环境变量设置）
ADMIN_USERNAME = os.environ.get('DRIVER_MANAGER_USERNAME', 'admin')
ADMIN_PASSWORD = os.environ.get('DRIVER_MANAGER_PASSWORD', 'admin')
ALLOW_PRIVATE_ONLY = os.environ.get('DRIVER_MANAGER_ALLOW_PRIVATE_ONLY', '1').lower() not in ('0', 'false', 'no', 'off')

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['MAX_CONTENT_LENGTH'] = MAX_CONTENT_LENGTH

# 确保上传目录存在
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

@app.before_request
def limit_private_network():
    """默认只允许本机、链路本地和内网地址访问。"""
    if not ALLOW_PRIVATE_ONLY:
        return None

    remote_addr = request.remote_addr or ''
    try:
        addr = ipaddress.ip_address(remote_addr)
    except ValueError:
        return Response('仅允许内网访问驱动管理器\n', 403)

    if addr.is_loopback or addr.is_private or addr.is_link_local:
        return None

    return Response('仅允许内网访问驱动管理器\n', 403)

def _is_within_directory(base_dir, target_path):
    base = os.path.realpath(base_dir)
    target = os.path.realpath(target_path)
    return os.path.commonpath([base, target]) == base

def safe_extract_tar(tar, path):
    """防止恶意 tar 包写出解压目录。"""
    for member in tar.getmembers():
        target_path = os.path.join(path, member.name)
        if not _is_within_directory(path, target_path):
            raise ValueError(f'压缩包包含非法路径: {member.name}')
        if member.issym() or member.islnk():
            raise ValueError(f'压缩包包含链接文件: {member.name}')
    tar.extractall(path)

def safe_extract_zip(zip_ref, path):
    """防止恶意 zip 包写出解压目录。"""
    for member in zip_ref.infolist():
        target_path = os.path.join(path, member.filename)
        if not _is_within_directory(path, target_path):
            raise ValueError(f'压缩包包含非法路径: {member.filename}')
    zip_ref.extractall(path)

def check_auth(username, password):
    """验证用户名和密码"""
    return username == ADMIN_USERNAME and password == ADMIN_PASSWORD

def authenticate():
    """发送401响应"""
    return Response(
        '需要登录才能访问驱动管理器\n',
        401,
        {'WWW-Authenticate': 'Basic realm="CUPS Driver Manager"'}
    )

def requires_auth(f):
    """认证装饰器"""
    @wraps(f)
    def decorated(*args, **kwargs):
        auth = request.authorization
        if not auth or not check_auth(auth.username, auth.password):
            return authenticate()
        return f(*args, **kwargs)
    return decorated

def allowed_file(filename):
    """检查文件类型是否允许"""
    if '.' not in filename:
        return False
    file_type = get_file_type(filename)
    if file_type == 'unknown':
        return False
    ext = filename.rsplit('.', 1)[1].lower()
    return ext in ALLOWED_EXTENSIONS or file_type in ('tar.gz', 'ppd')

def get_file_type(filename):
    """获取文件类型"""
    filename_lower = filename.lower()
    if filename_lower.endswith('.deb'):
        return 'deb'
    elif filename_lower.endswith('.ppd') or filename_lower.endswith('.ppd.gz'):
        return 'ppd'
    elif filename_lower.endswith('.tar.gz') or filename_lower.endswith('.tgz'):
        return 'tar.gz'
    elif filename_lower.endswith('.tar'):
        return 'tar'
    elif filename_lower.endswith('.zip'):
        return 'zip'
    elif filename_lower.endswith('.rpm'):
        return 'rpm'
    elif filename_lower.endswith('.sh') or filename_lower.endswith('.run'):
        return 'script'
    else:
        return 'unknown'

def upload_extension(filename):
    """保留真实扩展名，避免中文文件名被 secure_filename 清空后丢失类型。"""
    filename_lower = filename.lower()
    for ext in ('.tar.gz', '.ppd.gz', '.tgz', '.deb', '.ppd', '.tar', '.zip', '.rpm', '.sh', '.run'):
        if filename_lower.endswith(ext):
            return ext
    return Path(filename_lower).suffix

def safe_upload_filename(filename):
    safe_name = secure_filename(filename)
    ext = upload_extension(filename)
    if not safe_name:
        safe_name = f'driver{ext}'
    elif ext and safe_name.lower() == ext.lstrip('.'):
        safe_name = f'driver{ext}'
    elif ext and not safe_name.lower().endswith(ext):
        safe_name = f'{Path(safe_name).stem or "driver"}{ext}'
    return f'{uuid.uuid4().hex}_{safe_name}'

def find_files_by_suffix(root_dir, suffixes):
    suffixes = tuple(s.lower() for s in suffixes)
    matches = []
    for path in Path(root_dir).rglob('*'):
        if path.is_file() and path.name.lower().endswith(suffixes):
            matches.append(path)
    return sorted(matches, key=lambda p: str(p).lower())

def run_command(cmd, shell=False, cwd=None, input_text=None):
    """执行命令并返回结果"""
    try:
        if shell:
            result = subprocess.run(
                cmd, shell=True, capture_output=True, text=True,
                timeout=COMMAND_TIMEOUT, cwd=cwd, input=input_text,
                encoding='utf-8', errors='replace'
            )
        else:
            result = subprocess.run(
                cmd, capture_output=True, text=True, timeout=COMMAND_TIMEOUT,
                cwd=cwd, input=input_text, encoding='utf-8', errors='replace'
            )
        return {
            'success': result.returncode == 0,
            'stdout': result.stdout,
            'stderr': result.stderr,
            'returncode': result.returncode
        }
    except subprocess.TimeoutExpired:
        return {'success': False, 'stdout': '', 'stderr': f'命令执行超时（>{COMMAND_TIMEOUT}秒）', 'returncode': -1}
    except Exception as e:
        return {'success': False, 'stdout': '', 'stderr': str(e), 'returncode': -1}

def install_deb(filepath):
    """安装 .deb 包"""
    results = []

    result = run_command(['dpkg', '-i', filepath])
    if result['success']:
        results.append(('安装 DEB 包', result))
        return results

    # dpkg 因依赖缺失失败是常见情况，继续修复依赖后重试，不把初次失败计为最终失败。
    results.append(('初次安装 DEB 包', {
        'success': True,
        'stdout': result['stdout'],
        'stderr': '初次 dpkg 安装未完成，正在尝试自动修复依赖后重试。\n' + result['stderr'],
        'returncode': 0
    }))
    fix_result = run_command(['apt-get', 'install', '-f', '-y'])
    results.append(('修复 DEB 依赖', fix_result))
    if fix_result['success']:
        retry_result = run_command(['dpkg', '-i', filepath])
        results.append(('重新安装 DEB 包', retry_result))

    return results

def install_ppd(filepath):
    """安装 .ppd 文件"""
    results = []

    # PPD文件目录
    ppd_dirs = [
        '/usr/share/ppd/custom',
        '/usr/share/cups/model'
    ]

    # 确保目录存在
    for ppd_dir in ppd_dirs:
        os.makedirs(ppd_dir, exist_ok=True)

    filename = os.path.basename(filepath)
    is_gzipped = filename.lower().endswith('.ppd.gz')
    dest_filename = filename[:-3] if is_gzipped else filename

    for ppd_dir in ppd_dirs:
        dest = os.path.join(ppd_dir, dest_filename)
        try:
            if is_gzipped:
                with gzip.open(filepath, 'rb') as src, open(dest, 'wb') as dst:
                    shutil.copyfileobj(src, dst)
            else:
                shutil.copy2(filepath, dest)
            os.chmod(dest, 0o644)
            results.append((f'安装 PPD 到 {ppd_dir}', {
                'success': True,
                'stdout': f'已安装到 {dest}',
                'stderr': '',
                'returncode': 0
            }))
        except Exception as e:
            results.append((f'安装 PPD 到 {ppd_dir}', {
                'success': False,
                'stdout': '',
                'stderr': str(e),
                'returncode': 1
            }))

    return results

def install_extracted_dir(extract_dir):
    """从已解压目录中自动寻找可安装的驱动文件。"""
    results = []
    root_path = Path(extract_dir)

    deb_files = find_files_by_suffix(extract_dir, ('.deb',))
    if deb_files:
        for deb_file in deb_files:
            results.extend(install_deb(str(deb_file)))
        return results

    rpm_files = find_files_by_suffix(extract_dir, ('.rpm',))
    if rpm_files:
        for rpm_file in rpm_files:
            results.extend(install_rpm(str(rpm_file)))
        return results

    install_script_names = {'install.sh', 'setup.sh', 'install', 'setup'}
    for path in root_path.rglob('*'):
        if path.is_file() and path.name.lower() in install_script_names:
            os.chmod(path, 0o755)
            if path.name.lower().endswith('.sh'):
                result = run_command(['/bin/sh', str(path)], cwd=str(path.parent))
            else:
                result = run_command([str(path)], cwd=str(path.parent))
            results.append(('执行安装脚本', result))
            return results

    for path in root_path.rglob('*'):
        if path.is_file() and path.name.lower() == 'makefile':
            make_result = run_command(['make', '-C', str(path.parent)])
            results.append(('make', make_result))
            if make_result['success']:
                install_result = run_command(['make', '-C', str(path.parent), 'install'])
                results.append(('make install', install_result))
            return results

    ppd_files = find_files_by_suffix(extract_dir, ('.ppd', '.ppd.gz'))
    if ppd_files:
        for ppd_file in ppd_files:
            results.extend(install_ppd(str(ppd_file)))
        return results

    results.append(('查找安装方式', {
        'success': False,
        'stdout': '',
        'stderr': '未找到 DEB、RPM、安装脚本、Makefile 或 PPD 文件',
        'returncode': 1
    }))
    return results

def install_tar_gz(filepath):
    """安装 .tar/.tar.gz/.tgz 包"""
    results = []
    extract_dir = tempfile.mkdtemp(prefix='driver_')

    try:
        with tarfile.open(filepath, 'r:*') as tar:
            safe_extract_tar(tar, extract_dir)
        results.append(('解压文件', {
            'success': True,
            'stdout': f'已解压到 {extract_dir}',
            'stderr': '',
            'returncode': 0
        }))
        results.extend(install_extracted_dir(extract_dir))
    except Exception as e:
        results.append(('解压文件', {
            'success': False,
            'stdout': '',
            'stderr': str(e),
            'returncode': 1
        }))
    finally:
        # 清理临时目录
        shutil.rmtree(extract_dir, ignore_errors=True)

    return results

def install_zip(filepath):
    """安装 .zip 包"""
    results = []

    # 创建临时解压目录
    extract_dir = tempfile.mkdtemp(prefix='driver_')

    try:
        # 解压
        with zipfile.ZipFile(filepath, 'r') as zip_ref:
            safe_extract_zip(zip_ref, extract_dir)
        results.append(('解压文件', {
            'success': True,
            'stdout': f'已解压到 {extract_dir}',
            'stderr': '',
            'returncode': 0
        }))
        results.extend(install_extracted_dir(extract_dir))

    except Exception as e:
        results.append(('解压文件', {
            'success': False,
            'stdout': '',
            'stderr': str(e),
            'returncode': 1
        }))
    finally:
        # 清理临时目录
        shutil.rmtree(extract_dir, ignore_errors=True)

    return results

def install_rpm(filepath):
    """安装 .rpm 包（转换为deb后安装）"""
    results = []

    # 检查 alien 是否安装
    alien_check = run_command(['which', 'alien'])
    if not alien_check['success']:
        # 安装 alien
        install_result = run_command(['apt-get', 'install', '-y', 'alien'])
        results.append(('安装alien工具', install_result))
        if not install_result['success']:
            return results

    # 使用 alien 转换
    work_dir = os.path.dirname(filepath)
    convert_result = run_command(['alien', '-d', filepath], cwd=work_dir)
    results.append(('转换RPM为DEB', convert_result))

    if convert_result['success']:
        # 查找生成的 deb 文件
        deb_files = list(Path(work_dir).glob('*.deb'))
        for deb_file in deb_files:
            deb_results = install_deb(str(deb_file))
            results.extend(deb_results)
            os.remove(deb_file)  # 清理临时deb文件

    return results

def install_script(filepath):
    """执行安装脚本"""
    results = []
    filename = os.path.basename(filepath).lower()

    os.chmod(filepath, 0o755)

    # HP 插件需要特殊处理（非交互式安装）
    if 'hplip' in filename and 'plugin' in filename:
        # 直接执行 .run 文件，使用 yes 自动确认交互式提示
        # 注意：hp-plugin 命令需要 .asc 签名文件，比较麻烦，所以直接执行
        result = run_command(['/bin/sh', filepath], input_text='y\n' * 20)
        results.append(('执行 HP 插件安装脚本', result))
    else:
        # 普通脚本直接执行，避免文件名包含空格时 shell 拼接失败。
        if filename.endswith('.sh'):
            result = run_command(['/bin/sh', filepath])
        else:
            result = run_command([filepath])
        results.append(('执行安装脚本', result))

    return results

def install_driver(filepath, file_type):
    """根据文件类型安装驱动"""
    if file_type == 'deb':
        return install_deb(filepath)
    elif file_type == 'ppd':
        return install_ppd(filepath)
    elif file_type in ('tar.gz', 'tar', 'tgz'):
        return install_tar_gz(filepath)
    elif file_type == 'zip':
        return install_zip(filepath)
    elif file_type == 'rpm':
        return install_rpm(filepath)
    elif file_type == 'script':
        return install_script(filepath)
    else:
        return [('未知类型', {
            'success': False,
            'stdout': '',
            'stderr': f'不支持的文件类型: {file_type}',
            'returncode': 1
        })]

# HTML模板
HTML_TEMPLATE = '''
<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>CUPS 驱动管理器</title>
    <style>
        * {
            box-sizing: border-box;
            margin: 0;
            padding: 0;
        }
        body {
            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif;
            background: #f5f5f5;
            min-height: 100vh;
        }
        .header {
            background: rgba(46,46,46,.9);
            color: white;
            padding: 10px 20px;
            position: fixed;
            top: 0;
            left: 0;
            right: 0;
            z-index: 100;
        }
        .header ul {
            list-style: none;
            display: flex;
            gap: 5px;
        }
        .header a {
            color: white;
            text-decoration: none;
            padding: 5px 10px;
            display: block;
        }
        .header a:hover, .header a.active {
            background: rgba(255,255,255,0.1);
            border-radius: 3px;
        }
        .container {
            max-width: 900px;
            margin: 0 auto;
            padding: 80px 20px 40px;
        }
        h1 {
            color: #333;
            margin-bottom: 20px;
            border-bottom: 2px solid #007bff;
            padding-bottom: 10px;
        }
        .card {
            background: white;
            border-radius: 8px;
            box-shadow: 0 2px 10px rgba(0,0,0,0.1);
            padding: 25px;
            margin-bottom: 20px;
        }
        .card h2 {
            color: #333;
            margin-bottom: 15px;
            font-size: 1.2em;
        }
        .upload-area {
            border: 2px dashed #ccc;
            border-radius: 8px;
            padding: 40px;
            text-align: center;
            cursor: pointer;
            transition: all 0.3s;
            margin-bottom: 15px;
        }
        .upload-area:hover, .upload-area.dragover {
            border-color: #007bff;
            background: #f0f7ff;
        }
        .upload-area input[type="file"] {
            display: none;
        }
        .upload-area .icon {
            font-size: 48px;
            color: #ccc;
            margin-bottom: 10px;
        }
        .btn {
            background: #007bff;
            color: white;
            border: none;
            padding: 12px 24px;
            border-radius: 5px;
            cursor: pointer;
            font-size: 16px;
            transition: background 0.3s;
        }
        .btn:hover {
            background: #0056b3;
        }
        .btn:disabled {
            background: #ccc;
            cursor: not-allowed;
        }
        .btn-secondary {
            background: #6c757d;
        }
        .btn-secondary:hover {
            background: #545b62;
        }
        .file-info {
            background: #e9ecef;
            padding: 10px 15px;
            border-radius: 5px;
            margin: 15px 0;
            display: none;
        }
        .file-info.show {
            display: block;
        }
        .supported-types {
            color: #666;
            font-size: 14px;
            margin-top: 10px;
        }
        .supported-types span {
            background: #e9ecef;
            padding: 2px 8px;
            border-radius: 3px;
            margin: 2px;
            display: inline-block;
        }
        .text-muted {
            color: #666;
            font-size: 13px;
            margin-top: 8px;
        }
        .alert {
            padding: 15px;
            border-radius: 5px;
            margin-bottom: 15px;
        }
        .alert-success {
            background: #d4edda;
            color: #155724;
            border: 1px solid #c3e6cb;
        }
        .alert-danger {
            background: #f8d7da;
            color: #721c24;
            border: 1px solid #f5c6cb;
        }
        .alert-info {
            background: #d1ecf1;
            color: #0c5460;
            border: 1px solid #bee5eb;
        }
        .result-box {
            background: #1e1e1e;
            color: #d4d4d4;
            padding: 15px;
            border-radius: 5px;
            font-family: 'Consolas', 'Monaco', monospace;
            font-size: 13px;
            max-height: 400px;
            overflow-y: auto;
            white-space: pre-wrap;
            word-wrap: break-word;
        }
        .result-item {
            margin-bottom: 15px;
            padding-bottom: 15px;
            border-bottom: 1px solid #333;
        }
        .result-item:last-child {
            border-bottom: none;
            margin-bottom: 0;
            padding-bottom: 0;
        }
        .result-item .step {
            color: #569cd6;
            font-weight: bold;
        }
        .result-item .success {
            color: #4ec9b0;
        }
        .result-item .error {
            color: #f14c4c;
        }
        .loading {
            display: none;
            text-align: center;
            padding: 20px;
        }
        .loading.show {
            display: block;
        }
        .spinner {
            border: 3px solid #f3f3f3;
            border-top: 3px solid #007bff;
            border-radius: 50%;
            width: 30px;
            height: 30px;
            animation: spin 1s linear infinite;
            margin: 0 auto 10px;
        }
        @keyframes spin {
            0% { transform: rotate(0deg); }
            100% { transform: rotate(360deg); }
        }
        .driver-list {
            list-style: none;
        }
        .driver-list li {
            padding: 10px;
            border-bottom: 1px solid #eee;
            display: flex;
            justify-content: space-between;
            align-items: center;
        }
        .driver-list li:last-child {
            border-bottom: none;
        }
        .quick-links {
            display: flex;
            gap: 10px;
            flex-wrap: wrap;
            margin-top: 15px;
        }
        .quick-links a {
            padding: 8px 15px;
            background: #e9ecef;
            color: #333;
            text-decoration: none;
            border-radius: 5px;
            font-size: 14px;
        }
        .quick-links a:hover {
            background: #dee2e6;
        }
    </style>
</head>
<body>
    <div class="header">
        <ul>
            <li><a href="http://{{ request.host.split(':')[0] }}:631/" target="_blank">CUPS 管理</a></li>
            <li><a href="/" class="active">驱动管理器</a></li>
            <li><a href="/drivers">已安装驱动</a></li>
        </ul>
    </div>

    <div class="container">
        <h1>CUPS 打印机驱动管理器</h1>

        {% with messages = get_flashed_messages(with_categories=true) %}
            {% if messages %}
                {% for category, message in messages %}
                    <div class="alert alert-{{ category }}">{{ message }}</div>
                {% endfor %}
            {% endif %}
        {% endwith %}

        <div class="card">
            <h2>上传并安装驱动</h2>
            <form id="uploadForm" action="/upload" method="post" enctype="multipart/form-data">
                <div class="upload-area" id="uploadArea">
                    <div class="icon">📦</div>
                    <p>点击或拖拽文件到此处上传</p>
                    <input type="file" name="driver_file" id="fileInput" accept=".deb,.ppd,.tar.gz,.tgz,.tar,.zip,.rpm,.sh,.run">
                </div>
                <div class="file-info" id="fileInfo">
                    <strong>已选择：</strong><span id="fileName"></span>
                    <span id="fileSize"></span>
                </div>
                <div class="supported-types">
                    支持的格式：
                    <span>.deb</span>
                    <span>.ppd</span>
                    <span>.ppd.gz</span>
                    <span>.tar</span>
                    <span>.tgz</span>
                    <span>.tar.gz</span>
                    <span>.zip</span>
                    <span>.rpm</span>
                    <span>.sh</span>
                    <span>.run</span>
                </div>
                <p class="text-muted">最大上传：{{ max_upload_mb }} MB。压缩包内可自动识别 DEB、RPM、PPD、PPD.GZ、安装脚本和 Makefile。</p>
                <br>
                <button type="submit" class="btn" id="uploadBtn" disabled>开始安装</button>
            </form>

            <div class="loading" id="loading">
                <div class="spinner"></div>
                <p>正在安装驱动，请稍候...</p>
            </div>
        </div>

        {% if results %}
        <div class="card">
            <h2>安装结果</h2>
            <div class="result-box">
                {% for step, result in results %}
                <div class="result-item">
                    <div class="step">▶ {{ step }}</div>
                    {% if result.success %}
                    <div class="success">✓ 成功</div>
                    {% else %}
                    <div class="error">✗ 失败 (返回码: {{ result.returncode }})</div>
                    {% endif %}
                    {% if result.stdout %}<div>{{ result.stdout }}</div>{% endif %}
                    {% if result.stderr %}<div class="error">{{ result.stderr }}</div>{% endif %}
                </div>
                {% endfor %}
            </div>
        </div>
        {% endif %}

        <div class="card">
            <h2>快速链接</h2>
            <div class="quick-links">
                <a href="http://{{ request.host.split(':')[0] }}:631/admin" target="_blank">CUPS 管理页面</a>
                <a href="http://{{ request.host.split(':')[0] }}:631/printers" target="_blank">打印机列表</a>
                <a href="/drivers">查看已安装驱动</a>
            </div>
        </div>
    </div>

    <script>
        const uploadArea = document.getElementById('uploadArea');
        const fileInput = document.getElementById('fileInput');
        const fileInfo = document.getElementById('fileInfo');
        const fileName = document.getElementById('fileName');
        const fileSize = document.getElementById('fileSize');
        const uploadBtn = document.getElementById('uploadBtn');
        const uploadForm = document.getElementById('uploadForm');
        const loading = document.getElementById('loading');

        uploadArea.addEventListener('click', () => fileInput.click());

        uploadArea.addEventListener('dragover', (e) => {
            e.preventDefault();
            uploadArea.classList.add('dragover');
        });

        uploadArea.addEventListener('dragleave', () => {
            uploadArea.classList.remove('dragover');
        });

        uploadArea.addEventListener('drop', (e) => {
            e.preventDefault();
            uploadArea.classList.remove('dragover');
            if (e.dataTransfer.files.length) {
                fileInput.files = e.dataTransfer.files;
                updateFileInfo();
            }
        });

        fileInput.addEventListener('change', updateFileInfo);

        function updateFileInfo() {
            if (fileInput.files.length) {
                const file = fileInput.files[0];
                fileName.textContent = file.name;
                fileSize.textContent = ' (' + formatSize(file.size) + ')';
                fileInfo.classList.add('show');
                uploadBtn.disabled = false;
            }
        }

        function formatSize(bytes) {
            if (bytes < 1024) return bytes + ' B';
            if (bytes < 1024 * 1024) return (bytes / 1024).toFixed(1) + ' KB';
            return (bytes / 1024 / 1024).toFixed(1) + ' MB';
        }

        uploadForm.addEventListener('submit', () => {
            uploadBtn.disabled = true;
            loading.classList.add('show');
        });
    </script>
</body>
</html>
'''

DRIVERS_TEMPLATE = '''
<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>已安装驱动 - CUPS 驱动管理器</title>
    <style>
        * { box-sizing: border-box; margin: 0; padding: 0; }
        body {
            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Arial, sans-serif;
            background: #f5f5f5;
            min-height: 100vh;
        }
        .header {
            background: rgba(46,46,46,.9);
            color: white;
            padding: 10px 20px;
            position: fixed;
            top: 0;
            left: 0;
            right: 0;
            z-index: 100;
        }
        .header ul { list-style: none; display: flex; gap: 5px; }
        .header a {
            color: white;
            text-decoration: none;
            padding: 5px 10px;
            display: block;
        }
        .header a:hover, .header a.active {
            background: rgba(255,255,255,0.1);
            border-radius: 3px;
        }
        .container {
            max-width: 900px;
            margin: 0 auto;
            padding: 80px 20px 40px;
        }
        h1 {
            color: #333;
            margin-bottom: 20px;
            border-bottom: 2px solid #007bff;
            padding-bottom: 10px;
        }
        .card {
            background: white;
            border-radius: 8px;
            box-shadow: 0 2px 10px rgba(0,0,0,0.1);
            padding: 25px;
            margin-bottom: 20px;
        }
        .card h2 {
            color: #333;
            margin-bottom: 15px;
            font-size: 1.2em;
        }
        .driver-list { list-style: none; }
        .driver-list li {
            padding: 8px 0;
            border-bottom: 1px solid #eee;
            font-family: monospace;
            font-size: 13px;
        }
        .driver-list li:last-child { border-bottom: none; }
        .section { margin-bottom: 30px; }
        .section h3 {
            color: #555;
            margin-bottom: 10px;
            font-size: 1em;
        }
        .empty { color: #999; font-style: italic; }
        .btn {
            background: #007bff;
            color: white;
            border: none;
            padding: 8px 16px;
            border-radius: 5px;
            cursor: pointer;
            text-decoration: none;
            display: inline-block;
            margin-top: 15px;
        }
        .btn:hover { background: #0056b3; }
    </style>
</head>
<body>
    <div class="header">
        <ul>
            <li><a href="http://{{ request.host.split(':')[0] }}:631/" target="_blank">CUPS 管理</a></li>
            <li><a href="/">驱动管理器</a></li>
            <li><a href="/drivers" class="active">已安装驱动</a></li>
        </ul>
    </div>

    <div class="container">
        <h1>已安装的打印机驱动</h1>

        <div class="card">
            <div class="section">
                <h3>PPD 文件 (/usr/share/ppd/custom/)</h3>
                {% if ppd_custom %}
                <ul class="driver-list">
                    {% for ppd in ppd_custom %}
                    <li>{{ ppd }}</li>
                    {% endfor %}
                </ul>
                {% else %}
                <p class="empty">暂无自定义PPD文件</p>
                {% endif %}
            </div>

            <div class="section">
                <h3>CUPS 模型 (/usr/share/cups/model/)</h3>
                {% if cups_model %}
                <ul class="driver-list">
                    {% for model in cups_model %}
                    <li>{{ model }}</li>
                    {% endfor %}
                </ul>
                {% else %}
                <p class="empty">暂无自定义模型文件</p>
                {% endif %}
            </div>

            <div class="section">
                <h3>已安装的打印机相关软件包</h3>
                {% if packages %}
                <ul class="driver-list">
                    {% for pkg in packages %}
                    <li>{{ pkg }}</li>
                    {% endfor %}
                </ul>
                {% else %}
                <p class="empty">未找到打印机相关软件包</p>
                {% endif %}
            </div>
        </div>

        <a href="/" class="btn">← 返回上传页面</a>
    </div>
</body>
</html>
'''

@app.route('/')
@requires_auth
def index():
    return render_template_string(HTML_TEMPLATE, results=None, max_upload_mb=MAX_UPLOAD_MB)

@app.errorhandler(413)
def upload_too_large(error):
    if request.path.startswith('/api/'):
        return jsonify({
            'success': False,
            'error': f'文件过大，当前最大允许上传 {MAX_UPLOAD_MB} MB'
        }), 413
    flash(f'文件过大，当前最大允许上传 {MAX_UPLOAD_MB} MB', 'danger')
    return redirect(url_for('index'))

@app.route('/upload', methods=['POST'])
@requires_auth
def upload_file():
    if 'driver_file' not in request.files:
        flash('没有选择文件', 'danger')
        return redirect(url_for('index'))

    file = request.files['driver_file']

    if file.filename == '':
        flash('没有选择文件', 'danger')
        return redirect(url_for('index'))

    if not allowed_file(file.filename):
        flash('不支持的文件类型', 'danger')
        return redirect(url_for('index'))

    # 保存文件
    filename = safe_upload_filename(file.filename)
    filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    file.save(filepath)

    # 获取文件类型
    file_type = get_file_type(filename)

    # 安装驱动
    results = install_driver(filepath, file_type)

    # 清理上传的文件
    try:
        os.remove(filepath)
    except:
        pass

    # 检查是否全部成功
    all_success = all(r[1]['success'] for r in results)
    if all_success:
        flash('驱动安装成功！', 'success')
    else:
        flash('驱动安装过程中出现错误，请查看详细信息', 'danger')

    return render_template_string(HTML_TEMPLATE, results=results, max_upload_mb=MAX_UPLOAD_MB)

@app.route('/drivers')
@requires_auth
def list_drivers():
    # 获取自定义PPD文件
    ppd_custom = []
    ppd_dir = '/usr/share/ppd/custom'
    if os.path.exists(ppd_dir):
        ppd_custom = os.listdir(ppd_dir)

    # 获取CUPS模型
    cups_model = []
    model_dir = '/usr/share/cups/model'
    if os.path.exists(model_dir):
        for f in os.listdir(model_dir):
            if f.endswith('.ppd') or f.endswith('.ppd.gz'):
                cups_model.append(f)

    # 获取已安装的打印机相关软件包
    packages = []
    result = run_command(['dpkg', '-l'])
    if result['success']:
        for line in result['stdout'].split('\n'):
            if any(keyword in line.lower() for keyword in ['printer', 'cups', 'hplip', 'gutenprint', 'foomatic', 'epson', 'canon', 'brother', 'samsung', 'pantum']):
                parts = line.split()
                if len(parts) >= 3 and parts[0] == 'ii':
                    packages.append(f"{parts[1]} ({parts[2]})")

    return render_template_string(DRIVERS_TEMPLATE,
                                  ppd_custom=ppd_custom,
                                  cups_model=cups_model,
                                  packages=packages[:50])  # 限制显示数量

@app.route('/api/install', methods=['POST'])
@requires_auth
def api_install():
    """API接口，返回JSON"""
    if 'driver_file' not in request.files:
        return jsonify({'success': False, 'error': '没有选择文件'})

    file = request.files['driver_file']

    if file.filename == '':
        return jsonify({'success': False, 'error': '没有选择文件'})

    if not allowed_file(file.filename):
        return jsonify({'success': False, 'error': '不支持的文件类型'})

    # 保存文件
    filename = safe_upload_filename(file.filename)
    filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    file.save(filepath)

    # 获取文件类型并安装
    file_type = get_file_type(filename)
    results = install_driver(filepath, file_type)

    # 清理
    try:
        os.remove(filepath)
    except:
        pass

    all_success = all(r[1]['success'] for r in results)
    return jsonify({
        'success': all_success,
        'results': [{'step': r[0], 'success': r[1]['success'],
                     'stdout': r[1]['stdout'], 'stderr': r[1]['stderr']}
                    for r in results]
    })

def main():
    import argparse
    parser = argparse.ArgumentParser(description='CUPS 打印机驱动管理器')
    parser.add_argument('-p', '--port', type=int, default=632, help='监听端口 (默认: 632)')
    parser.add_argument('-H', '--host', default='0.0.0.0', help='监听地址 (默认: 0.0.0.0)')
    parser.add_argument('--username', default=None, help='管理员用户名')
    parser.add_argument('--password', default=None, help='管理员密码')
    args = parser.parse_args()

    global ADMIN_USERNAME, ADMIN_PASSWORD
    if args.username:
        ADMIN_USERNAME = args.username
    if args.password:
        ADMIN_PASSWORD = args.password

    print(f"CUPS 驱动管理器启动中...")
    print(f"访问地址: http://localhost:{args.port}/")
    print(f"用户名: {ADMIN_USERNAME}")
    print(f"密码: {ADMIN_PASSWORD}")
    print("-" * 40)

    app.run(host=args.host, port=args.port, debug=False)

if __name__ == '__main__':
    main()