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

import os
import sys
import subprocess
import tempfile
import shutil
import tarfile
import zipfile
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)

# 配置
UPLOAD_FOLDER = '/tmp/cups-drivers'
MAX_CONTENT_LENGTH = 100 * 1024 * 1024  # 100MB
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')

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

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

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
    ext = filename.rsplit('.', 1)[1].lower()
    # 处理 .tar.gz 情况
    if filename.endswith('.tar.gz'):
        return True
    return ext in ALLOWED_EXTENSIONS

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 run_command(cmd, shell=False):
    """执行命令并返回结果"""
    try:
        if shell:
            result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=300)
        else:
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        return {
            'success': result.returncode == 0,
            'stdout': result.stdout,
            'stderr': result.stderr,
            'returncode': result.returncode
        }
    except subprocess.TimeoutExpired:
        return {'success': False, 'stdout': '', 'stderr': '命令执行超时', '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])
    results.append(('dpkg -i', result))

    # 修复依赖
    if not result['success']:
        fix_result = run_command(['apt-get', 'install', '-f', '-y'])
        results.append(('apt-get install -f', fix_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)

    # 复制PPD文件
    filename = os.path.basename(filepath)
    dest = os.path.join(ppd_dirs[0], filename)

    try:
        shutil.copy2(filepath, dest)
        os.chmod(dest, 0o644)
        results.append(('复制PPD文件', {
            'success': True,
            'stdout': f'已复制到 {dest}',
            'stderr': '',
            'returncode': 0
        }))
    except Exception as e:
        results.append(('复制PPD文件', {
            'success': False,
            'stdout': '',
            'stderr': str(e),
            'returncode': 1
        }))

    return results

def install_tar_gz(filepath):
    """安装 .tar.gz 包"""
    results = []

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

    try:
        # 解压
        with tarfile.open(filepath, 'r:gz') as tar:
            tar.extractall(extract_dir)
        results.append(('解压文件', {
            'success': True,
            'stdout': f'已解压到 {extract_dir}',
            'stderr': '',
            'returncode': 0
        }))

        # 查找安装脚本
        install_scripts = ['install.sh', 'setup.sh', 'install', 'setup']
        found_script = None

        for root, dirs, files in os.walk(extract_dir):
            for script in install_scripts:
                if script in files:
                    found_script = os.path.join(root, script)
                    break
            if found_script:
                break

        if found_script:
            os.chmod(found_script, 0o755)
            result = run_command(found_script, shell=True)
            results.append(('执行安装脚本', result))
        else:
            # 查找 Makefile
            for root, dirs, files in os.walk(extract_dir):
                if 'Makefile' in files:
                    make_result = run_command(['make', '-C', root])
                    results.append(('make', make_result))
                    if make_result['success']:
                        install_result = run_command(['make', '-C', root, 'install'])
                        results.append(('make install', install_result))
                    break
            else:
                # 查找 PPD 文件
                ppd_files = list(Path(extract_dir).rglob('*.ppd'))
                if ppd_files:
                    for ppd_file in ppd_files:
                        ppd_results = install_ppd(str(ppd_file))
                        results.extend(ppd_results)
                else:
                    results.append(('查找安装方式', {
                        'success': False,
                        'stdout': '',
                        'stderr': '未找到安装脚本、Makefile或PPD文件',
                        'returncode': 1
                    }))

    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:
            zip_ref.extractall(extract_dir)
        results.append(('解压文件', {
            'success': True,
            'stdout': f'已解压到 {extract_dir}',
            'stderr': '',
            'returncode': 0
        }))

        # 查找 deb 文件
        deb_files = list(Path(extract_dir).rglob('*.deb'))
        if deb_files:
            for deb_file in deb_files:
                deb_results = install_deb(str(deb_file))
                results.extend(deb_results)
            return results

        # 查找安装脚本
        install_scripts = ['install.sh', 'setup.sh', 'install', 'setup']
        found_script = None

        for root, dirs, files in os.walk(extract_dir):
            for script in install_scripts:
                if script in files:
                    found_script = os.path.join(root, script)
                    break
            if found_script:
                break

        if found_script:
            os.chmod(found_script, 0o755)
            result = run_command(found_script, shell=True)
            results.append(('执行安装脚本', result))
        else:
            # 查找 PPD 文件
            ppd_files = list(Path(extract_dir).rglob('*.ppd'))
            if ppd_files:
                for ppd_file in ppd_files:
                    ppd_results = install_ppd(str(ppd_file))
                    results.extend(ppd_results)
            else:
                results.append(('查找安装方式', {
                    'success': False,
                    'stdout': '',
                    'stderr': '未找到deb包、安装脚本或PPD文件',
                    'returncode': 1
                }))

    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(f'cd {work_dir} && alien -d {filepath}', shell=True)
    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:
        # 方法1：尝试使用 hp-plugin 命令安装
        hp_plugin_check = run_command(['which', 'hp-plugin'])
        if hp_plugin_check['success']:
            result = run_command(f'hp-plugin -i -p {filepath}', shell=True)
            results.append(('使用 hp-plugin 安装 HP 插件', result))
        else:
            # 方法2：使用 yes 命令自动确认交互式提示
            result = run_command(f'yes | {filepath}', shell=True)
            results.append(('执行 HP 插件安装脚本', result))
    else:
        # 普通脚本直接执行
        result = run_command(filepath, shell=True)
        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;
        }
        .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>.tar.gz</span>
                    <span>.zip</span>
                    <span>.rpm</span>
                    <span>.sh</span>
                </div>
                <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)

@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 = secure_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)

@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 = secure_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('--password', default=None, help='管理员密码')
    args = parser.parse_args()

    global ADMIN_PASSWORD
    if args.password:
        ADMIN_PASSWORD = args.password

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

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

if __name__ == '__main__':
    main()