ZeroSend/src/transfer/webdav.rs

use crate::config::SharedFolderConfig;
use axum::{
    body::Body,
    extract::{Path, State},
    http::{header, Method, Request, Response, StatusCode},
    routing::any,
    Router,
};
use std::sync::Arc;
use tokio::fs::File;
use tokio::sync::RwLock;
use tokio_util::io::ReaderStream;
use tracing::{info, warn};

pub struct WebDavState {
    pub shared_folders: Arc<RwLock<Vec<SharedFolderConfig>>>,
    pub port: u16,
}

pub fn create_webdav_router(state: Arc<WebDavState>) -> Router {
    Router::new()
        .route("/webdav", any(handle_webdav_root))
        .route("/webdav/{*path}", any(handle_webdav_path))
        .with_state(state)
}

pub async fn start_webdav_server(
    shared_folders: Arc<RwLock<Vec<SharedFolderConfig>>>,
    port: u16,
) {
    let state = Arc::new(WebDavState {
        shared_folders,
        port,
    });
    let app = create_webdav_router(state);
    let addr = std::net::SocketAddr::from(([127, 0, 0, 1], port));

    info!("Starting local WebDAV server on http://{}", addr);

    match tokio::net::TcpListener::bind(addr).await {
        Ok(listener) => {
            let _ = axum::serve(listener, app).await;
        }
        Err(e) => {
            warn!("Failed to bind WebDAV server on {}: {}", addr, e);
        }
    }
}

async fn handle_webdav_root(
    State(state): State<Arc<WebDavState>>,
    req: Request<Body>,
) -> Response<Body> {
    let method = req.method().clone();
    let headers = req.headers().clone();

    if method == Method::OPTIONS {
        return options_response();
    }

    if method.as_str() == "PROPFIND" {
        let shares = state.shared_folders.read().await;
        let depth = headers
            .get("Depth")
            .and_then(|h| h.to_str().ok())
            .unwrap_or("1");

        let mut xml = String::from("<?xml version=\"1.0\" encoding=\"utf-8\"?>\n<D:multistatus xmlns:D=\"DAV:\">\n");

        // Root response
        xml.push_str("  <D:response>\n");
        xml.push_str("    <D:href>/webdav/</D:href>\n");
        xml.push_str("    <D:propstat>\n");
        xml.push_str("      <D:prop>\n");
        xml.push_str("        <D:displayname>ZeroSend WebDAV</D:displayname>\n");
        xml.push_str("        <D:resourcetype><D:collection/></D:resourcetype>\n");
        xml.push_str("      </D:prop>\n");
        xml.push_str("      <D:status>HTTP/1.1 200 OK</D:status>\n");
        xml.push_str("    </D:propstat>\n");
        xml.push_str("  </D:response>\n");

        if depth != "0" {
            for s in shares.iter() {
                let encoded_name = urlencoding::encode(&s.name);
                xml.push_str("  <D:response>\n");
                xml.push_str(&format!("    <D:href>/webdav/{}/</D:href>\n", encoded_name));
                xml.push_str("    <D:propstat>\n");
                xml.push_str("      <D:prop>\n");
                xml.push_str(&format!("        <D:displayname>{}</D:displayname>\n", escape_xml(&s.name)));
                xml.push_str("        <D:resourcetype><D:collection/></D:resourcetype>\n");
                xml.push_str("      </D:prop>\n");
                xml.push_str("      <D:status>HTTP/1.1 200 OK</D:status>\n");
                xml.push_str("    </D:propstat>\n");
                xml.push_str("  </D:response>\n");
            }
        }

        xml.push_str("</D:multistatus>");

        return Response::builder()
            .status(StatusCode::MULTI_STATUS)
            .header(header::CONTENT_TYPE, "application/xml; charset=utf-8")
            .header("DAV", "1, 2")
            .header("MS-Author-Via", "DAV")
            .body(Body::from(xml))
            .unwrap();
    }

    options_response()
}

async fn handle_webdav_path(
    State(state): State<Arc<WebDavState>>,
    Path(raw_path): Path<String>,
    req: Request<Body>,
) -> Response<Body> {
    let method = req.method().clone();
    let headers = req.headers().clone();

    if method == Method::OPTIONS {
        return options_response();
    }

    let decoded_path = urlencoding::decode(&raw_path)
        .unwrap_or(std::borrow::Cow::Borrowed(&raw_path))
        .to_string();
    let parts: Vec<&str> = decoded_path.split('/').filter(|s| !s.is_empty()).collect();

    if parts.is_empty() {
        return handle_webdav_root(State(state), req).await;
    }

    let share_name = parts[0];
    let subpath_parts = &parts[1..];

    let shares = state.shared_folders.read().await;
    let folder = match shares.iter().find(|s| s.name == share_name || s.id == share_name) {
        Some(f) => f.clone(),
        None => {
            return Response::builder()
                .status(StatusCode::NOT_FOUND)
                .body(Body::from("Folder not found"))
                .unwrap();
        }
    };
    drop(shares);

    let base_canonical = match folder.path.canonicalize() {
        Ok(p) => p,
        Err(_) => {
            return Response::builder()
                .status(StatusCode::NOT_FOUND)
                .body(Body::from("Base path not found"))
                .unwrap();
        }
    };

    let mut target_path = base_canonical.clone();
    for seg in subpath_parts {
        if *seg == ".." || *seg == "." {
            continue;
        }
        target_path.push(seg);
    }

    // Safety check against path traversal
    if !target_path.starts_with(&base_canonical) {
        return Response::builder()
            .status(StatusCode::FORBIDDEN)
            .body(Body::from("Forbidden path"))
            .unwrap();
    }

    match method.as_str() {
        "PROPFIND" => {
            let depth = headers
                .get("Depth")
                .and_then(|h| h.to_str().ok())
                .unwrap_or("1");

            let mut xml = String::from("<?xml version=\"1.0\" encoding=\"utf-8\"?>\n<D:multistatus xmlns:D=\"DAV:\">\n");

            if let Ok(meta) = std::fs::metadata(&target_path) {
                let is_dir = meta.is_dir();
                let href_path = format!("/webdav/{}{}", raw_path.trim_end_matches('/'), if is_dir { "/" } else { "" });

                xml.push_str("  <D:response>\n");
                xml.push_str(&format!("    <D:href>{}</D:href>\n", href_path));
                xml.push_str("    <D:propstat>\n");
                xml.push_str("      <D:prop>\n");
                let name = target_path.file_name().map(|n| n.to_string_lossy().into_owned()).unwrap_or_else(|| share_name.to_string());
                xml.push_str(&format!("        <D:displayname>{}</D:displayname>\n", escape_xml(&name)));
                if is_dir {
                    xml.push_str("        <D:resourcetype><D:collection/></D:resourcetype>\n");
                } else {
                    xml.push_str("        <D:resourcetype/>\n");
                    xml.push_str(&format!("        <D:getcontentlength>{}</D:getcontentlength>\n", meta.len()));
                    let mime = crate::transfer::get_mime_type(&name);
                    xml.push_str(&format!("        <D:getcontenttype>{}</D:getcontenttype>\n", mime));
                }
                xml.push_str("      </D:prop>\n");
                xml.push_str("      <D:status>HTTP/1.1 200 OK</D:status>\n");
                xml.push_str("    </D:propstat>\n");
                xml.push_str("  </D:response>\n");

                if is_dir && depth != "0" {
                    if let Ok(entries) = std::fs::read_dir(&target_path) {
                        for entry in entries.flatten() {
                            if let Ok(child_meta) = entry.metadata() {
                                let child_name = entry.file_name().to_string_lossy().into_owned();
                                let child_is_dir = child_meta.is_dir();
                                let child_href = format!("{}/{}{}", href_path.trim_end_matches('/'), urlencoding::encode(&child_name), if child_is_dir { "/" } else { "" });

                                xml.push_str("  <D:response>\n");
                                xml.push_str(&format!("    <D:href>{}</D:href>\n", child_href));
                                xml.push_str("    <D:propstat>\n");
                                xml.push_str("      <D:prop>\n");
                                xml.push_str(&format!("        <D:displayname>{}</D:displayname>\n", escape_xml(&child_name)));
                                if child_is_dir {
                                    xml.push_str("        <D:resourcetype><D:collection/></D:resourcetype>\n");
                                } else {
                                    xml.push_str("        <D:resourcetype/>\n");
                                    xml.push_str(&format!("        <D:getcontentlength>{}</D:getcontentlength>\n", child_meta.len()));
                                    let mime = crate::transfer::get_mime_type(&child_name);
                                    xml.push_str(&format!("        <D:getcontenttype>{}</D:getcontenttype>\n", mime));
                                }
                                xml.push_str("      </D:prop>\n");
                                xml.push_str("      <D:status>HTTP/1.1 200 OK</D:status>\n");
                                xml.push_str("    </D:propstat>\n");
                                xml.push_str("  </D:response>\n");
                            }
                        }
                    }
                }
            } else {
                return Response::builder()
                    .status(StatusCode::NOT_FOUND)
                    .body(Body::from("Path not found"))
                    .unwrap();
            }

            xml.push_str("</D:multistatus>");

            Response::builder()
                .status(StatusCode::MULTI_STATUS)
                .header(header::CONTENT_TYPE, "application/xml; charset=utf-8")
                .header("DAV", "1, 2")
                .header("MS-Author-Via", "DAV")
                .body(Body::from(xml))
                .unwrap()
        }
        "GET" | "HEAD" => {
            if target_path.is_dir() {
                return Response::builder()
                    .status(StatusCode::OK)
                    .header(header::CONTENT_TYPE, "text/html; charset=utf-8")
                    .body(Body::from("<h1>Directory Index</h1>"))
                    .unwrap();
            }

            match File::open(&target_path).await {
                Ok(file) => {
                    let meta = file.metadata().await.ok();
                    let file_size = meta.as_ref().map(|m| m.len()).unwrap_or(0);
                    let filename = target_path.file_name().map(|f| f.to_string_lossy().into_owned()).unwrap_or_default();
                    let mime = crate::transfer::get_mime_type(&filename);

                    let stream = ReaderStream::new(file);
                    Response::builder()
                        .status(StatusCode::OK)
                        .header(header::CONTENT_TYPE, mime)
                        .header(header::CONTENT_LENGTH, file_size.to_string())
                        .header(header::ACCEPT_RANGES, "bytes")
                        .body(Body::from_stream(stream))
                        .unwrap()
                }
                Err(_) => Response::builder()
                    .status(StatusCode::NOT_FOUND)
                    .body(Body::from("File not found"))
                    .unwrap(),
            }
        }
        "PUT" => {
            if folder.read_only {
                return Response::builder()
                    .status(StatusCode::FORBIDDEN)
                    .body(Body::from("Folder is read-only"))
                    .unwrap();
            }

            let bytes = match axum::body::to_bytes(req.into_body(), 100 * 1024 * 1024).await {
                Ok(b) => b,
                Err(e) => {
                    return Response::builder()
                        .status(StatusCode::BAD_REQUEST)
                        .body(Body::from(format!("Failed to read body: {}", e)))
                        .unwrap();
                }
            };

            if let Some(parent) = target_path.parent() {
                let _ = std::fs::create_dir_all(parent);
            }

            if std::fs::write(&target_path, bytes).is_ok() {
                Response::builder()
                    .status(StatusCode::CREATED)
                    .body(Body::empty())
                    .unwrap()
            } else {
                Response::builder()
                    .status(StatusCode::INTERNAL_SERVER_ERROR)
                    .body(Body::from("Failed to write file"))
                    .unwrap()
            }
        }
        "MKCOL" => {
            if folder.read_only {
                return Response::builder()
                    .status(StatusCode::FORBIDDEN)
                    .body(Body::from("Folder is read-only"))
                    .unwrap();
            }

            if std::fs::create_dir_all(&target_path).is_ok() {
                Response::builder()
                    .status(StatusCode::CREATED)
                    .body(Body::empty())
                    .unwrap()
            } else {
                Response::builder()
                    .status(StatusCode::METHOD_NOT_ALLOWED)
                    .body(Body::from("Failed to create folder"))
                    .unwrap()
            }
        }
        "DELETE" => {
            if folder.read_only {
                return Response::builder()
                    .status(StatusCode::FORBIDDEN)
                    .body(Body::from("Folder is read-only"))
                    .unwrap();
            }

            if target_path.is_dir() {
                let _ = std::fs::remove_dir_all(&target_path);
            } else {
                let _ = std::fs::remove_file(&target_path);
            }

            Response::builder()
                .status(StatusCode::NO_CONTENT)
                .body(Body::empty())
                .unwrap()
        }
        _ => options_response(),
    }
}

fn options_response() -> Response<Body> {
    Response::builder()
        .status(StatusCode::OK)
        .header("DAV", "1, 2")
        .header("MS-Author-Via", "DAV")
        .header("Allow", "OPTIONS, GET, HEAD, POST, DELETE, PROPFIND, PROPPATCH, COPY, MOVE, MKCOL, PUT, LOCK, UNLOCK")
        .body(Body::empty())
        .unwrap()
}

fn escape_xml(s: &str) -> String {
    s.replace('&', "&amp;")
        .replace('<', "&lt;")
        .replace('>', "&gt;")
        .replace('"', "&quot;")
        .replace('\'', "&apos;")
}

mod urlencoding {
    pub fn encode(s: &str) -> String {
        url::form_urlencoded::byte_serialize(s.as_bytes()).collect()
    }

    pub fn decode(s: &str) -> Result<std::borrow::Cow<'_, str>, ()> {
        url::form_urlencoded::parse(s.as_bytes())
            .next()
            .map(|(k, _)| k)
            .ok_or(())
    }
}

/// Helper function to mount Windows Network Drive via `net use`
pub fn mount_webdav_drive(drive_letter: char, port: u16) -> Result<String, String> {
    let drive = format!("{}:", drive_letter);
    let url = format!("http://127.0.0.1:{}/webdav", port);

    // Unmount first if already mapped
    let _ = std::process::Command::new("net")
        .args(["use", &drive, "/delete", "/y"])
        .output();

    let output = std::process::Command::new("net")
        .args(["use", &drive, &url])
        .output()
        .map_err(|e| format!("Failed to run net use: {}", e))?;

    if output.status.success() {
        Ok(format!("Successfully mounted {} to {}", drive, url))
    } else {
        let err = String::from_utf8_lossy(&output.stderr);
        Err(format!("Mount failed: {}", err.trim()))
    }
}

/// Helper function to unmount Windows Network Drive
pub fn unmount_webdav_drive(drive_letter: char) -> Result<(), String> {
    let drive = format!("{}:", drive_letter);
    let output = std::process::Command::new("net")
        .args(["use", &drive, "/delete", "/y"])
        .output()
        .map_err(|e| format!("Failed to run net use: {}", e))?;

    if output.status.success() {
        Ok(())
    } else {
        let err = String::from_utf8_lossy(&output.stderr);
        Err(format!("Unmount failed: {}", err.trim()))
    }
}