fileshare-app/src/ui.rs

use crate::{models::*, net, storage};
use axum::{
    body::{Body, Bytes},
    extract::{Path, Request, State},
    http::{header, HeaderMap, HeaderName, HeaderValue, Method, StatusCode},
    middleware::Next,
    response::{Html, IntoResponse, Response},
    Json,
};
use futures_util::StreamExt;
use std::{
    io,
    net::IpAddr,
    sync::Arc,
    time::{Duration, Instant},
};
use tokio_stream::wrappers::ReceiverStream;
use tokio_util::io::ReaderStream;

const INDEX_HTML: &str = include_str!("static/index.html");
const APP_JS: &str = include_str!("static/app.js");

const CSP: &str = "default-src 'none'; script-src 'self'; style-src 'unsafe-inline'; img-src 'self'; \
                   connect-src 'self'; base-uri 'none'; form-action 'none'; frame-ancestors 'none'";

// ───────────────────────── Защита UI (CSRF / DNS-rebinding) ─────────────────────────
pub struct UiGuard {
    pub token: String,
    pub cookie_name: String,
    pub hosts: Vec<String>,
}

fn ct_eq(a: &[u8], b: &[u8]) -> bool {
    a.len() == b.len() && a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0
}

fn cookie_ok(headers: &HeaderMap, g: &UiGuard) -> bool {
    headers
        .get_all(header::COOKIE)
        .iter()
        .filter_map(|v| v.to_str().ok())
        .flat_map(|s| s.split(';'))
        .any(|kv| {
            kv.trim()
                .strip_prefix(g.cookie_name.as_str())
                .and_then(|r| r.strip_prefix('='))
                .map_or(false, |v| ct_eq(v.as_bytes(), g.token.as_bytes()))
        })
}

fn with_security_headers(mut resp: Response) -> Response {
    let items: [(HeaderName, &'static str); 5] = [
        (header::CONTENT_SECURITY_POLICY, CSP),
        (HeaderName::from_static("x-content-type-options"), "nosniff"),
        (header::X_FRAME_OPTIONS, "DENY"),
        (header::REFERRER_POLICY, "no-referrer"),
        (header::CACHE_CONTROL, "no-store"),
    ];
    let h = resp.headers_mut();
    for (k, v) in items {
        h.entry(k).or_insert(HeaderValue::from_static(v)); // не перетираем заголовки обработчика
    }
    resp
}

/// 1) Host только 127.0.0.1/localhost:порт (анти-rebinding)
/// 2) POST только с Origin нашего же хоста (анти-CSRF)
/// 3) Доступ по токену из консоли: /?t=… выставляет HttpOnly SameSite=Strict cookie
pub async fn ui_guard(State(g): State<Arc<UiGuard>>, req: Request, next: Next) -> Response {
    let headers = req.headers();

    let host_ok = headers
        .get(header::HOST)
        .and_then(|v| v.to_str().ok())
        .map_or(false, |h| g.hosts.iter().any(|a| a.eq_ignore_ascii_case(h)));
    if !host_ok {
        return StatusCode::FORBIDDEN.into_response();
    }

    if !matches!(*req.method(), Method::GET | Method::HEAD) {
        let origin_ok = headers
            .get(header::ORIGIN)
            .and_then(|v| v.to_str().ok())
            .and_then(|o| o.strip_prefix("http://"))
            .map_or(false, |o| g.hosts.iter().any(|a| a.eq_ignore_ascii_case(o)));
        if !origin_ok {
            return StatusCode::FORBIDDEN.into_response();
        }
    }

    if req.uri().path() == "/" {
        if let Some(q) = req.uri().query() {
            let token_ok = q
                .split('&')
                .filter_map(|kv| kv.strip_prefix("t="))
                .any(|t| ct_eq(t.as_bytes(), g.token.as_bytes()));
            if token_ok {
                let mut resp = StatusCode::SEE_OTHER.into_response();
                resp.headers_mut().insert(header::LOCATION, HeaderValue::from_static("/"));
                if let Ok(c) = HeaderValue::from_str(&format!(
                    "{}={}; Path=/; HttpOnly; SameSite=Strict",
                    g.cookie_name, g.token
                )) {
                    resp.headers_mut().insert(header::SET_COOKIE, c);
                }
                return with_security_headers(resp);
            }
        }
    }

    if !cookie_ok(headers, &g) {
        return with_security_headers(
            (StatusCode::FORBIDDEN, "Откройте адрес из консоли (с токеном ?t=…)").into_response(),
        );
    }
    with_security_headers(next.run(req).await)
}

// ───────────────────────── Статика и состояние ─────────────────────────
pub async fn ui_index() -> Html<&'static str> {
    Html(INDEX_HTML)
}

pub async fn ui_js() -> impl IntoResponse {
    ([(header::CONTENT_TYPE, "text/javascript; charset=utf-8")], APP_JS)
}

pub async fn ui_get_state(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
    let mut peers: Vec<PeerDevice> = state.peers.read().values().cloned().collect();
    peers.sort_by(|a, b| a.name.cmp(&b.name));
    let texts: Vec<TextItem> = state.ram_texts.read().iter().cloned().collect();
    let media: Vec<MediaItem> = state.disk_media.read().iter().cloned().collect();
    let transfers: Vec<TransferProgress> = state.transfers.read().clone();
    let scan = state.scan_status.read().clone();

    Json(serde_json::json!({
        "device_name": state.device_name,
        "my_ip": state.my_ip,
        "p2p_port": state.p2p_port,
        "my_fingerprint": state.my_fingerprint,
        "my_pin": short_id(&state.my_fingerprint),
        "media_dir": state.media_dir.display().to_string(),
        "peers": peers,
        "texts": texts,
        "media": media,
        "transfers": transfers,
        "scan": scan,
    }))
}

// ───────────────────────── Устройства ─────────────────────────
pub async fn ui_add_manual_peer(
    State(state): State<Arc<AppState>>,
    Json(req): Json<UiManualPeerRequest>,
) -> Result<StatusCode, AppError> {
    let ip: IpAddr = req
        .ip
        .trim()
        .parse()
        .map_err(|_| AppError::new(StatusCode::BAD_REQUEST, "Некорректный IP-адрес"))?;
    match ip {
        IpAddr::V4(v4) if is_lan_v4(v4) => {}
        _ => return Err(AppError::new(StatusCode::BAD_REQUEST, "Разрешены только локальные IPv4-адреса")),
    }
    let port = req.port.filter(|p| *p != 0).unwrap_or(state.p2p_port);
    net::probe_and_register(&state, ip, port)
        .await
        .map_err(|e| AppError::new(StatusCode::BAD_GATEWAY, format!("Узел {}: {}", ip, e)))?;
    Ok(StatusCode::OK)
}

pub async fn ui_scan_range(
    State(state): State<Arc<AppState>>,
    Json(req): Json<UiScanRangeRequest>,
) -> Result<StatusCode, AppError> {
    let port = req.port.filter(|p| *p != 0).unwrap_or(state.p2p_port);
    let (label, ips) =
        parse_ip_range(&req.range, &state.my_ip).map_err(|e| AppError::new(StatusCode::BAD_REQUEST, e))?;
    let total = ips.len();

    {
        // проверка и установка флага — атомарно, под одним write-замком
        let mut s = state.scan_status.write();
        if s.active {
            return Err(AppError::new(StatusCode::CONFLICT, "Сканирование уже выполняется"));
        }
        *s = ScanProgress {
            active: true,
            range: label.clone(),
            scanned: 0,
            total,
            found: 0,
        };
    }
    state.add_log(format!("Запущено сканирование {} ({} адресов, порт {})", label, total, port));

    let st = state.clone();
    let worker = tokio::spawn(net::run_scan(state.clone(), ips, port));
    tokio::spawn(async move {
        let _ = worker.await; // даже при панике воркера флаг active будет сброшен
        let found = {
            let mut s = st.scan_status.write();
            s.active = false;
            s.found
        };
        st.add_log(format!("Сканирование {} завершено. Найдено устройств: {}", label, found));
    });
    Ok(StatusCode::OK)
}

pub async fn ui_ping_peers(State(state): State<Arc<AppState>>) -> StatusCode {
    tokio::spawn(async move { net::ping_all_peers(&state).await });
    StatusCode::OK
}

pub async fn ui_trust_peer(
    State(state): State<Arc<AppState>>,
    Json(req): Json<UiTrustRequest>,
) -> Result<StatusCode, AppError> {
    let msg = {
        let mut peers = state.peers.write();
        let p = peers
            .get_mut(&req.peer_id)
            .ok_or_else(|| AppError::new(StatusCode::NOT_FOUND, "Устройство не найдено"))?;
        p.verified = req.trust;
        if req.trust {
            format!("Устройство '{}' [{}] ЗАВЕРЕНО", p.name, p.pin_code)
        } else {
            format!("Доверие устройству '{}' [{}] ОТОЗВАНО", p.name, p.pin_code)
        }
    };
    state.persist_trust();
    state.add_log(msg);
    Ok(StatusCode::OK)
}

pub async fn ui_clear_ram(State(state): State<Arc<AppState>>) -> StatusCode {
    state.ram_texts.write().clear();
    state.add_log("Текстовый буфер в ОЗУ очищен");
    StatusCode::OK
}

// ───────────────────────── Отправка ─────────────────────────
fn resolve_target(state: &AppState, id: &str) -> Result<(Vec<IpAddr>, u16, String), AppError> {
    let peers = state.peers.read();
    let p = peers
        .get(id)
        .ok_or_else(|| AppError::new(StatusCode::NOT_FOUND, "Получатель не найден"))?;
    if !p.verified {
        return Err(AppError::new(
            StatusCode::FORBIDDEN,
            format!("Сначала заверьте устройство '{}' (код сверки {})", p.name, p.pin_code),
        ));
    }
    let ips: Vec<IpAddr> = p.ips.iter().filter_map(|s| s.parse().ok()).collect();
    if ips.is_empty() {
        return Err(AppError::new(
            StatusCode::BAD_REQUEST,
            "Адрес получателя неизвестен. Выполните опрос или сканирование.",
        ));
    }
    Ok((ips, p.port, p.name.clone()))
}

pub async fn ui_send_text(
    State(state): State<Arc<AppState>>,
    Json(req): Json<UiSendTextRequest>,
) -> Result<StatusCode, AppError> {
    if req.text.trim().is_empty() {
        return Err(AppError::new(StatusCode::BAD_REQUEST, "Пустой текст"));
    }
    if req.text.len() > MAX_TEXT_BYTES {
        return Err(AppError::new(StatusCode::PAYLOAD_TOO_LARGE, "Текст больше 64 КБ"));
    }
    let (ips, port, peer_name) = resolve_target(&state, &req.peer_id)?;
    let client = state.client_for_peer(&req.peer_id, false)?;
    let body = serde_json::to_vec(&IncomingTextPayload { text: req.text })
        .map_err(|e| AppError::internal("json", e))?;

    for ip in ips {
        let url = format!("https://{}:{}/p2p/text", ip, port);
        match client
            .post(&url)
            .header(header::CONTENT_TYPE, "application/json")
            .body(body.clone())
            .send()
            .await
        {
            Ok(r) if r.status().is_success() => {
                state.add_log(format!("Текст доставлен в ОЗУ устройства '{}'", peer_name));
                return Ok(StatusCode::OK);
            }
            Ok(r) if r.status() == StatusCode::FORBIDDEN => {
                return Err(AppError::new(
                    StatusCode::FORBIDDEN,
                    "Получатель ещё не заверил ваше устройство (доверие должно быть взаимным)",
                ))
            }
            Ok(r) => {
                return Err(AppError::new(
                    StatusCode::BAD_GATEWAY,
                    format!("Получатель отклонил запрос ({})", r.status()),
                ))
            }
            Err(_) => continue, // пробуем следующий адрес
        }
    }
    Err(AppError::new(
        StatusCode::BAD_GATEWAY,
        "Не удалось установить защищённое соединение (узел недоступен или сертификат не совпал)",
    ))
}

/// Потоковая отправка: браузер → (канал с backpressure) → TLS → получатель. Память O(1).
pub async fn ui_send_media(
    State(state): State<Arc<AppState>>,
    headers: HeaderMap,
    body: Body,
) -> Result<StatusCode, AppError> {
    let peer_id = hdr(&headers, "x-target-id")
        .ok_or_else(|| AppError::new(StatusCode::BAD_REQUEST, "Не указан получатель"))?
        .to_string();
    let display_name = sanitize_filename(&pct_decode(hdr(&headers, "x-filename").unwrap_or("file.bin")));
    let total: u64 = hdr(&headers, "content-length")
        .and_then(|v| v.parse().ok())
        .ok_or_else(|| AppError::new(StatusCode::LENGTH_REQUIRED, "Не указан размер файла"))?;
    if total == 0 {
        return Err(AppError::new(StatusCode::BAD_REQUEST, "Файл пустой"));
    }
    if total > MAX_FILE_SIZE {
        return Err(AppError::new(StatusCode::PAYLOAD_TOO_LARGE, "Файл больше 2 ГБ"));
    }

    let (ips, port, peer_name) = resolve_target(&state, &peer_id)?;
    let client = state.client_for_peer(&peer_id, true)?;
    let url = format!("https://{}:{}/p2p/media", ips[0], port);

    let tid = format!("tx_{}", random_hex(6));
    state.add_transfer(TransferProgress {
        id: tid.clone(),
        filename: display_name.clone(),
        peer_name: peer_name.clone(),
        direction: Direction::Out,
        transferred_bytes: 0,
        total_bytes: total,
        percent: 0,
        done: false,
        error: None,
        updated_at: now_sec(),
    });

    let (tx, rx) = tokio::sync::mpsc::channel::<Result<Bytes, io::Error>>(4);
    let fwd_state = state.clone();
    let fwd_id = tid.clone();
    let fwd = tokio::spawn(async move {
        let mut stream = std::pin::pin!(body.into_data_stream());
        let mut sent: u64 = 0;
        let mut last = Instant::now();
        while let Some(item) = stream.next().await {
            let chunk = match item {
                Ok(c) => c,
                Err(_) => {
                    let _ = tx.send(Err(io::Error::new(io::ErrorKind::Other, "client aborted"))).await;
                    return;
                }
            };
            sent += chunk.len() as u64;
            if sent > total {
                let _ = tx.send(Err(io::Error::new(io::ErrorKind::Other, "body too long"))).await;
                return;
            }
            // если получатель не читает 30 с — закрываем поток, запрос завершится ошибкой
            match tokio::time::timeout(Duration::from_secs(30), tx.send(Ok(chunk))).await {
                Ok(Ok(())) => {}
                _ => return,
            }
            if last.elapsed() >= Duration::from_millis(100) {
                fwd_state.update_transfer_bytes(&fwd_id, sent);
                last = Instant::now();
            }
        }
    });

    let result = client
        .post(&url)
        .header("x-filename", pct_encode(&display_name))
        .header("x-total-size", total.to_string())
        .header(header::CONTENT_TYPE, "application/octet-stream")
        .header(header::CONTENT_LENGTH, total.to_string())
        .body(reqwest::Body::wrap_stream(ReceiverStream::new(rx)))
        .send()
        .await;
    fwd.abort();

    let fail = |msg: &str, status: StatusCode| {
        state.fail_transfer(&tid, msg.to_string());
        remove_transfer_later(&state, tid.clone(), 10);
        AppError::new(status, msg)
    };

    match result {
        Ok(r) if r.status().is_success() => {
            state.finish_transfer(&tid);
            remove_transfer_later(&state, tid.clone(), 4);
            state.add_log(format!(
                "Файл '{}' ({} байт) доставлен на '{}'",
                display_name, total, peer_name
            ));
            Ok(StatusCode::OK)
        }
        Ok(r) if r.status() == StatusCode::FORBIDDEN => Err(fail(
            "Получатель ещё не заверил ваше устройство",
            StatusCode::FORBIDDEN,
        )),
        Ok(r) => Err(fail(
            &format!("Получатель отклонил передачу (код {})", r.status()),
            StatusCode::BAD_GATEWAY,
        )),
        Err(_) => Err(fail(
            "Ошибка сети/TLS при передаче (узел недоступен, сертификат не совпал или обрыв)",
            StatusCode::BAD_GATEWAY,
        )),
    }
}

// ───────────────────────── Раздача принятых файлов ─────────────────────────
pub async fn ui_serve_saved_file(State(state): State<Arc<AppState>>, Path(filename): Path<String>) -> Response {
    let not_found = || (StatusCode::NOT_FOUND, "Файл не найден").into_response();
    if filename.is_empty()
        || filename.starts_with('.')
        || filename.contains(|c: char| matches!(c, '/' | '\\' | '\0'))
    {
        return not_found();
    }
    let path = state.media_dir.join(&filename);
    let Ok(meta) = tokio::fs::metadata(&path).await else { return not_found() };
    if !meta.is_file() {
        return not_found();
    }
    let Ok(file) = tokio::fs::File::open(&path).await else { return not_found() };

    let inline = storage::is_previewable(&filename);
    // Content-Type определяем САМИ по расширению, а не по заявлению отправителя
    let ctype = if inline { storage::mime_from_name(&filename) } else { "application/octet-stream" };
    let disposition = format!(
        "{}; filename*=UTF-8''{}",
        if inline { "inline" } else { "attachment" },
        pct_encode(&filename)
    );

    let mut resp = Response::new(Body::from_stream(ReaderStream::new(file)));
    let h = resp.headers_mut();
    h.insert(header::CONTENT_TYPE, HeaderValue::from_static(ctype));
    h.insert(header::CONTENT_LENGTH, HeaderValue::from(meta.len()));
    if let Ok(v) = HeaderValue::from_str(&disposition) {
        h.insert(header::CONTENT_DISPOSITION, v);
    }
    h.insert(
        header::CONTENT_SECURITY_POLICY,
        HeaderValue::from_static("default-src 'none'; sandbox"),
    );
    h.insert(
        HeaderName::from_static("cross-origin-resource-policy"),
        HeaderValue::from_static("same-origin"),
    );
    resp
}