app.py

from __future__ import annotations

import mimetypes
import os
import socket
from datetime import datetime
from pathlib import Path
from typing import List, Dict, Any
import base64
import io

from flask import Flask, render_template, request, redirect, url_for, send_from_directory, flash, session, send_file, jsonify
from werkzeug.utils import secure_filename
import qrcode
from flask_socketio import SocketIO, emit
import zipfile
import io
from zeroconf import ServiceInfo, Zeroconf
import threading

BASE_DIR = Path(__file__).parent
SHARE_DIR = BASE_DIR / "shared"
SHARE_DIR.mkdir(parents=True, exist_ok=True)

app = Flask(__name__)
app.config["SECRET_KEY"] = "change-me"  

socketio = SocketIO(app)  


def human_size(num_bytes: int) -> str:
    """Convert a byte count into a readable size string."""
    step = 1024.0
    units = ["B", "KB", "MB", "GB", "TB"]
    size = float(num_bytes)
    for unit in units:
        if size < step:
            return f"{size:.1f} {unit}" if unit != "B" else f"{int(size)} {unit}"
        size /= step
    return f"{size:.1f} PB"


def list_files() -> List[Dict[str, Any]]:
    files = []
    for path in SHARE_DIR.iterdir():
        if path.is_file():
            stat = path.stat()
            files.append(
                {
                    "name": path.name,
                    "size": stat.st_size,
                    "size_display": human_size(stat.st_size),
                    "modified": datetime.fromtimestamp(stat.st_mtime),
                    "is_image": path.suffix.lower() in ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp'],
                    "is_text": path.suffix.lower() in ['.txt', '.md', '.py', '.js', '.html', '.css', '.json'],
                    "is_pdf": path.suffix.lower() == '.pdf',
                }
            )
    files.sort(key=lambda f: f["modified"], reverse=True)
    return files


def discover_local_ips() -> List[str]:
    """Return a list of non-loopback IPv4 addresses for display."""
    ips = set()

    # Primary interface guess via UDP connect trick
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.connect(("8.8.8.8", 80))
        primary_ip = sock.getsockname()[0]
        if primary_ip and not primary_ip.startswith("127."):
            ips.add(primary_ip)
    except OSError:
        pass
    finally:
        try:
            sock.close()
        except Exception:
            pass

    # Hostname resolution can reveal additional NICs
    try:
        for info in socket.getaddrinfo(socket.gethostname(), None):
            ip = info[4][0]
            if "." in ip and not ip.startswith("127."):
                ips.add(ip)
    except OSError:
        pass

    return sorted(ips)


def is_authenticated():
    view_password = os.environ.get("FTAPP_VIEW_PASSWORD")
    if not view_password:
        return True  # no password required
    return session.get('authenticated', False)


def require_auth():
    if not is_authenticated():
        return redirect(url_for('login'))


def generate_qr_code(url: str) -> str:
    """Generate a QR code for the given URL and return as base64 data URI."""
    qr = qrcode.QRCode(version=1, box_size=10, border=5)
    qr.add_data(url)
    qr.make(fit=True)
    img = qr.make_image(fill='black', back_color='white')
    buffer = io.BytesIO()
    img.save(buffer, format="PNG")
    buffer.seek(0)
    img_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
    return f"data:image/png;base64,{img_base64}"


@app.route("/login", methods=["GET", "POST"])
def login():
    if request.method == "POST":
        password = request.form.get("password")
        view_password = os.environ.get("FTAPP_VIEW_PASSWORD")
        if view_password and password == view_password:
            session['authenticated'] = True
            return redirect(url_for('home'))
        else:
            flash("Invalid password.", "error")
    return render_template("login.html")


@app.route("/", methods=["GET"])
def home():
    if not is_authenticated():
        return redirect(url_for('login'))
    files = list_files()
    ips = discover_local_ips()
    port = int(os.environ.get("FTAPP_PORT", "8000"))
    qr_codes = []
    for ip in ips:
        url = f"http://{ip}:{port}"
        qr_codes.append(generate_qr_code(url))
    return render_template("index.html", files=files, ips=ips, qr_codes=qr_codes, port=port)


@app.route("/upload", methods=["POST"])
def upload():
    upload_pin = os.environ.get("FTAPP_UPLOAD_PIN")
    if upload_pin:
        pin = request.form.get("pin")
        if pin != upload_pin:
            flash("Invalid upload PIN.", "error")
            return redirect(url_for("home"))

    files = request.files.getlist("files")
    valid_files = []

    for file in files:
        if not file or file.filename.strip() == "":
            continue
        filename = secure_filename(file.filename)
        if not filename:
            continue
        valid_files.append((filename, file))

    if not valid_files:
        flash("No valid files selected.", "error")
        return redirect(url_for("home"))

    saved = []
    for filename, file in valid_files:
        save_path = SHARE_DIR / filename
        file.save(save_path)
        saved.append(filename)

    if len(saved) == 1:
        message = f"Uploaded {saved[0]}"
    else:
        message = f"Uploaded {len(saved)} files"

    if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
        return jsonify({"status": "success", "message": message})
    else:
        flash(message, "success")
        socketio.emit('files_updated', {'action': 'upload', 'files': saved})
        return redirect(url_for("home"))


@app.route("/delete/<path:filename>", methods=["POST"])
def delete_file(filename: str):
    if not is_authenticated():
        return redirect(url_for('login'))
    file_path = SHARE_DIR / filename
    if file_path.exists() and file_path.is_file():
        file_path.unlink()
        flash(f"Deleted {filename}", "success")
        socketio.emit('files_updated', {'action': 'delete', 'files': [filename]})
    else:
        flash("File not found.", "error")
    return redirect(url_for("home"))


@app.route("/files/<path:filename>", methods=["GET"])
def download(filename: str):
    # send_from_directory will 404 on path traversal, keeping access scoped to SHARE_DIR
    guessed_type, _ = mimetypes.guess_type(filename)
    return send_from_directory(
        SHARE_DIR,
        filename,
        as_attachment=False,
        mimetype=guessed_type or "application/octet-stream",
    )


@app.route("/preview/<path:filename>", methods=["GET"])
def preview(filename: str):
    if not is_authenticated():
        return redirect(url_for('login'))
    file_path = SHARE_DIR / filename
    if not file_path.exists() or not file_path.is_file():
        return "File not found", 404
    suffix = file_path.suffix.lower()
    if suffix in ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp']:
        return send_from_directory(SHARE_DIR, filename, mimetype=mimetypes.guess_type(filename)[0])
    elif suffix == '.pdf':
        return send_from_directory(SHARE_DIR, filename, mimetype='application/pdf')
    elif suffix in ['.txt', '.md', '.py', '.js', '.html', '.css', '.json']:
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            return f"<pre>{content}</pre>", 200, {'Content-Type': 'text/html'}
        except UnicodeDecodeError:
            return "Cannot preview binary file", 400
    else:
        return "Preview not supported", 400


@app.route("/download_zip", methods=["POST"])
def download_zip():
    if not is_authenticated():
        return redirect(url_for('login'))
    selected_files = request.form.getlist("selected_files")
    if not selected_files:
        flash("No files selected.", "error")
        return redirect(url_for("home"))
    
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
        for filename in selected_files:
            file_path = SHARE_DIR / filename
            if file_path.exists() and file_path.is_file():
                zip_file.write(file_path, filename)
    
    zip_buffer.seek(0)
    return send_file(zip_buffer, as_attachment=True, download_name="files.zip", mimetype="application/zip")


@app.route("/health", methods=["GET"])
def health():
    return {"status": "ok", "files": len(list_files())}


def main():
    host = os.environ.get("FTAPP_HOST", "0.0.0.0")
    port = int(os.environ.get("FTAPP_PORT", "8000"))
    
    # Register mDNS service
    zeroconf = Zeroconf()
    service_info = ServiceInfo(
        "_http._tcp.local.",
        "localbeam._http._tcp.local.",
        addresses=[],  # Will be auto-detected
        port=port,
        properties={},
    )
    zeroconf.register_service(service_info)
    
    try:
        socketio.run(app, host=host, port=port, debug=False)
    finally:
        zeroconf.unregister_service(service_info)
        zeroconf.close()


if __name__ == "__main__":
    main()