266 lines
8.8 KiB
Python
266 lines
8.8 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""Pull a remote Windows folder via the win64_mp file API."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import base64
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
import urllib.error
|
||
|
|
import urllib.parse
|
||
|
|
import urllib.request
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
AUTH_USER = os.environ.get("AUTH_USER", "admin")
|
||
|
|
AUTH_PASS = os.environ.get("AUTH_PASS", "blueberries")
|
||
|
|
|
||
|
|
|
||
|
|
def win_join(parent: str, name: str) -> str:
|
||
|
|
parent = parent.rstrip("\\/")
|
||
|
|
name = name.replace("/", "\\").strip("\\")
|
||
|
|
if not name or name in (".", "..") or "\\" in name:
|
||
|
|
raise ValueError(f"bad entry name: {name!r}")
|
||
|
|
if parent.endswith(":"):
|
||
|
|
return parent + "\\" + name
|
||
|
|
return parent + "\\" + name
|
||
|
|
|
||
|
|
|
||
|
|
def norm_win(path: str) -> str:
|
||
|
|
return path.replace("/", "\\").rstrip("\\").lower()
|
||
|
|
|
||
|
|
|
||
|
|
def should_skip(remote: str, name: str, patterns: list[str]) -> bool:
|
||
|
|
remote_n = norm_win(remote)
|
||
|
|
name_n = name.replace("/", "\\").strip("\\").lower()
|
||
|
|
for raw in patterns:
|
||
|
|
pattern = norm_win(raw)
|
||
|
|
if not pattern:
|
||
|
|
continue
|
||
|
|
if "\\" not in pattern and ":" not in pattern:
|
||
|
|
if name_n == pattern:
|
||
|
|
return True
|
||
|
|
continue
|
||
|
|
if remote_n == pattern or remote_n.endswith("\\" + pattern) or remote_n.startswith(pattern + "\\"):
|
||
|
|
return True
|
||
|
|
return False
|
||
|
|
|
||
|
|
|
||
|
|
def split_skips(items: list[str]) -> list[str]:
|
||
|
|
out: list[str] = []
|
||
|
|
for item in items:
|
||
|
|
for part in item.split(","):
|
||
|
|
part = part.strip()
|
||
|
|
if part:
|
||
|
|
out.append(part)
|
||
|
|
return out
|
||
|
|
|
||
|
|
|
||
|
|
def prompt(label: str) -> str:
|
||
|
|
try:
|
||
|
|
value = input(f"{label}: ").strip()
|
||
|
|
except EOFError:
|
||
|
|
sys.exit("no input")
|
||
|
|
if not value:
|
||
|
|
sys.exit(f"{label} is required")
|
||
|
|
return value
|
||
|
|
|
||
|
|
|
||
|
|
def prompt_optional(label: str) -> str:
|
||
|
|
try:
|
||
|
|
return input(f"{label}: ").strip()
|
||
|
|
except EOFError:
|
||
|
|
return ""
|
||
|
|
|
||
|
|
|
||
|
|
def base_url(raw: str) -> str:
|
||
|
|
raw = raw.strip().rstrip("/")
|
||
|
|
if "://" not in raw:
|
||
|
|
raw = "http://" + raw
|
||
|
|
parsed = urllib.parse.urlparse(raw)
|
||
|
|
if not parsed.hostname:
|
||
|
|
sys.exit(f"bad endpoint: {raw}")
|
||
|
|
port = parsed.port or 5032
|
||
|
|
return f"{parsed.scheme}://{parsed.hostname}:{port}"
|
||
|
|
|
||
|
|
|
||
|
|
def local_root_for(win_path: str) -> Path:
|
||
|
|
name = win_path.rstrip("\\/").replace("/", "\\").rsplit("\\", 1)[-1]
|
||
|
|
if not name or name.endswith(":"):
|
||
|
|
name = "pull"
|
||
|
|
dest = Path("pulls") / name
|
||
|
|
dest.mkdir(parents=True, exist_ok=True)
|
||
|
|
return dest.resolve()
|
||
|
|
|
||
|
|
|
||
|
|
class Agent:
|
||
|
|
def __init__(self, origin: str) -> None:
|
||
|
|
self.origin = origin
|
||
|
|
self.auth = "Basic " + base64.b64encode(f"{AUTH_USER}:{AUTH_PASS}".encode()).decode()
|
||
|
|
|
||
|
|
def request(self, path: str, query: dict[str, str], dest: Path | None = None) -> object | None:
|
||
|
|
qs = urllib.parse.urlencode(query)
|
||
|
|
url = f"{self.origin}{path}"
|
||
|
|
if qs:
|
||
|
|
url += "?" + qs
|
||
|
|
req = urllib.request.Request(url, headers={"Authorization": self.auth})
|
||
|
|
try:
|
||
|
|
with urllib.request.urlopen(req, timeout=300) as resp:
|
||
|
|
if dest is None:
|
||
|
|
return json.load(resp)
|
||
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
with dest.open("wb") as out:
|
||
|
|
while True:
|
||
|
|
chunk = resp.read(1024 * 1024)
|
||
|
|
if not chunk:
|
||
|
|
break
|
||
|
|
out.write(chunk)
|
||
|
|
return None
|
||
|
|
except urllib.error.HTTPError as err:
|
||
|
|
detail = err.read().decode("utf-8", "replace")
|
||
|
|
try:
|
||
|
|
detail = json.loads(detail).get("error", detail)
|
||
|
|
except json.JSONDecodeError:
|
||
|
|
pass
|
||
|
|
raise RuntimeError(f"{err.code} {path}: {detail}") from err
|
||
|
|
except urllib.error.URLError as err:
|
||
|
|
raise RuntimeError(f"connect {self.origin}: {err.reason}") from err
|
||
|
|
|
||
|
|
def health(self) -> None:
|
||
|
|
body = self.request("/health", {})
|
||
|
|
if not isinstance(body, dict) or body.get("status") != "ok":
|
||
|
|
raise RuntimeError(f"unexpected health: {body}")
|
||
|
|
|
||
|
|
def list_dir(self, win_path: str) -> list[dict]:
|
||
|
|
body = self.request("/api/v1/files", {"path": win_path, "depth": "0"})
|
||
|
|
if not isinstance(body, dict):
|
||
|
|
raise RuntimeError(f"bad list response: {body}")
|
||
|
|
entries = body.get("entries") or []
|
||
|
|
if not isinstance(entries, list):
|
||
|
|
raise RuntimeError("list entries is not an array")
|
||
|
|
return entries
|
||
|
|
|
||
|
|
def download(self, win_path: str, dest: Path) -> None:
|
||
|
|
self.request("/api/v1/download", {"path": win_path}, dest=dest)
|
||
|
|
|
||
|
|
|
||
|
|
def already_have(local: Path, size: object) -> bool:
|
||
|
|
if not local.is_file():
|
||
|
|
return False
|
||
|
|
try:
|
||
|
|
return local.stat().st_size == int(size)
|
||
|
|
except (TypeError, ValueError):
|
||
|
|
return False
|
||
|
|
|
||
|
|
|
||
|
|
def pull(
|
||
|
|
agent: Agent, remote: str, dest_root: Path, rel: Path, skips: list[str]
|
||
|
|
) -> tuple[int, int, int]:
|
||
|
|
try:
|
||
|
|
entries = agent.list_dir(remote)
|
||
|
|
except RuntimeError as err:
|
||
|
|
print(f"skip dir {remote}: {err}", file=sys.stderr)
|
||
|
|
return 0, 0, 1
|
||
|
|
|
||
|
|
dest_root.joinpath(rel).mkdir(parents=True, exist_ok=True)
|
||
|
|
files_ok = 0
|
||
|
|
files_skip = 0
|
||
|
|
files_fail = 0
|
||
|
|
for entry in entries:
|
||
|
|
name = str(entry.get("name") or "")
|
||
|
|
kind = str(entry.get("type") or "")
|
||
|
|
try:
|
||
|
|
child = win_join(remote, name)
|
||
|
|
except ValueError as err:
|
||
|
|
print(f"skip {err}", file=sys.stderr)
|
||
|
|
files_fail += 1
|
||
|
|
continue
|
||
|
|
if should_skip(child, name, skips):
|
||
|
|
print(f"skip {child}")
|
||
|
|
files_skip += 1
|
||
|
|
continue
|
||
|
|
if kind == "dir":
|
||
|
|
ok, skipped, fail = pull(agent, child, dest_root, rel / name, skips)
|
||
|
|
files_ok += ok
|
||
|
|
files_skip += skipped
|
||
|
|
files_fail += fail
|
||
|
|
continue
|
||
|
|
if kind != "file":
|
||
|
|
print(f"skip {child}: unknown type {kind!r}", file=sys.stderr)
|
||
|
|
files_fail += 1
|
||
|
|
continue
|
||
|
|
local = dest_root / rel / name
|
||
|
|
if already_have(local, entry.get("size")):
|
||
|
|
files_skip += 1
|
||
|
|
continue
|
||
|
|
print(f"{child} ({entry.get('size', '?')} bytes)")
|
||
|
|
try:
|
||
|
|
agent.download(child, local)
|
||
|
|
files_ok += 1
|
||
|
|
except RuntimeError as err:
|
||
|
|
print(f" fail: {err}", file=sys.stderr)
|
||
|
|
files_fail += 1
|
||
|
|
return files_ok, files_skip, files_fail
|
||
|
|
|
||
|
|
|
||
|
|
def self_check() -> None:
|
||
|
|
assert win_join(r"C:\Users\me", "docs") == r"C:\Users\me\docs"
|
||
|
|
assert win_join(r"C:", "Windows") == r"C:\Windows"
|
||
|
|
assert base_url("10.0.0.5:5032") == "http://10.0.0.5:5032"
|
||
|
|
assert base_url("http://10.0.0.5") == "http://10.0.0.5:5032"
|
||
|
|
cache = r"C:\Users\Larisa\AppData\Local\BraveSoftware\Brave-Browser\User Data\Default\Service Worker\CacheStorage"
|
||
|
|
assert should_skip(cache, "CacheStorage", ["CacheStorage"])
|
||
|
|
assert should_skip(cache, "CacheStorage", [r"Default\Service Worker\CacheStorage"])
|
||
|
|
assert should_skip(cache, "CacheStorage", [cache])
|
||
|
|
assert not should_skip(r"C:\Users\Larisa\AppData\Local\BraveSoftware\Brave-Browser\User Data\Default", "Default", ["CacheStorage"])
|
||
|
|
assert split_skips(["CacheStorage,GPUCache", "Code Cache"]) == ["CacheStorage", "GPUCache", "Code Cache"]
|
||
|
|
print("ok")
|
||
|
|
|
||
|
|
|
||
|
|
def parse_args() -> argparse.Namespace:
|
||
|
|
parser = argparse.ArgumentParser(description="Download a remote Windows folder via win64_mp")
|
||
|
|
parser.add_argument("endpoint", nargs="?", help="ip:port of the agent")
|
||
|
|
parser.add_argument("remote", nargs="?", help="absolute Windows folder path")
|
||
|
|
parser.add_argument("dest", nargs="?", help="local destination directory")
|
||
|
|
parser.add_argument(
|
||
|
|
"--skip",
|
||
|
|
"-s",
|
||
|
|
action="append",
|
||
|
|
default=[],
|
||
|
|
metavar="NAME",
|
||
|
|
help="folder name or path to ignore; repeatable, comma-separated ok",
|
||
|
|
)
|
||
|
|
parser.add_argument("--self-check", action="store_true")
|
||
|
|
return parser.parse_args()
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
args = parse_args()
|
||
|
|
if args.self_check:
|
||
|
|
self_check()
|
||
|
|
return
|
||
|
|
endpoint = args.endpoint or prompt("Endpoint (ip:port)")
|
||
|
|
remote = args.remote or prompt("Windows folder path")
|
||
|
|
skips = split_skips(args.skip)
|
||
|
|
if not skips and args.endpoint is None:
|
||
|
|
skips = split_skips([prompt_optional("Skip folders (comma-separated, empty for none)")])
|
||
|
|
origin = base_url(endpoint)
|
||
|
|
dest = Path(args.dest).resolve() if args.dest else local_root_for(remote)
|
||
|
|
dest.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
agent = Agent(origin)
|
||
|
|
print(f"checking {origin}/health")
|
||
|
|
agent.health()
|
||
|
|
print(f"pulling {remote} -> {dest}")
|
||
|
|
if skips:
|
||
|
|
print("skipping: " + ", ".join(skips))
|
||
|
|
ok, skipped, fail = pull(agent, remote, dest, Path(), skips)
|
||
|
|
print(f"done: {ok} downloaded, {skipped} skipped, {fail} failures")
|
||
|
|
if fail:
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|