generation_destination.py 7.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203
  1. """Helpers for generation destinations and remote uploads."""
  2. from __future__ import annotations
  3. import shlex
  4. import subprocess
  5. import tempfile
  6. from dataclasses import dataclass
  7. from pathlib import Path
  8. from ..input import InputManager
  9. @dataclass
  10. class GenerationDestination:
  11. """Resolved generation target."""
  12. mode: str
  13. local_output_dir: Path | None = None
  14. remote_host: str | None = None
  15. remote_path: str | None = None
  16. @property
  17. def is_remote(self) -> bool:
  18. return self.mode == "remote"
  19. def normalize_output_path(path_value: str) -> Path:
  20. """Normalize paths that look absolute but were provided without a leading slash."""
  21. output_dir = Path(path_value)
  22. if not output_dir.is_absolute() and str(output_dir).startswith(("Users/", "home/", "usr/", "opt/", "var/", "tmp/")):
  23. output_dir = Path("/") / output_dir
  24. return output_dir
  25. def resolve_cli_destination(
  26. output: str | None,
  27. remote: str | None,
  28. remote_path: str | None,
  29. slug: str,
  30. ) -> GenerationDestination | None:
  31. """Resolve generation destination from explicit CLI flags only."""
  32. if output and remote:
  33. raise ValueError("Use either --output for a local directory or --remote for a remote server, not both")
  34. if remote_path and not remote:
  35. raise ValueError("--remote-path requires --remote")
  36. if remote:
  37. return GenerationDestination(
  38. mode="remote",
  39. remote_host=remote,
  40. remote_path=remote_path or f"~/{slug}",
  41. )
  42. if output:
  43. return GenerationDestination(mode="local", local_output_dir=normalize_output_path(output))
  44. return None
  45. def prompt_generation_destination(slug: str) -> GenerationDestination:
  46. """Prompt for local or remote generation target."""
  47. input_mgr = InputManager()
  48. destination_mode = input_mgr.numbered_choice("Store generated template in", ["local", "remote"], default="local")
  49. if destination_mode == "local":
  50. local_default = str(Path.cwd() / slug)
  51. local_output = input_mgr.text("Local output directory", default=local_default).strip() or local_default
  52. return GenerationDestination(mode="local", local_output_dir=normalize_output_path(local_output))
  53. remote_host = input_mgr.text("Remote server host or IP address", default=None).strip()
  54. if not remote_host:
  55. raise ValueError("Remote server host or IP address cannot be empty")
  56. remote_default = f"~/{slug}"
  57. remote_path = input_mgr.text("Remote target directory", default=remote_default).strip() or remote_default
  58. return GenerationDestination(mode="remote", remote_host=remote_host, remote_path=remote_path)
  59. def format_remote_destination(host: str, remote_path: str) -> str:
  60. """Format host/path for user-facing messages."""
  61. return f"{host}:{remote_path}"
  62. def build_remote_shell_path(remote_path: str, trailing_slash: bool = False) -> str:
  63. """Build a shell-safe remote path expression for ssh/scp commands."""
  64. normalized = remote_path.rstrip("/")
  65. suffix = "/" if trailing_slash else ""
  66. if normalized in {"~", ""}:
  67. return f'"$HOME"{suffix}'
  68. if normalized.startswith("~/"):
  69. relative = normalized[2:]
  70. quoted_relative = shlex.quote(f"{relative}{suffix}")
  71. return f'"$HOME"/{quoted_relative}'
  72. return shlex.quote(f"{normalized}{suffix}")
  73. def build_scp_remote_target(remote_host: str, remote_path: str, trailing_slash: bool = False) -> str:
  74. """Build an scp destination target for an already-resolved remote path."""
  75. normalized = remote_path.rstrip("/")
  76. suffix = "/" if trailing_slash else ""
  77. quoted_path = shlex.quote(f"{normalized}{suffix}")
  78. return f"{remote_host}:{quoted_path}"
  79. def resolve_remote_home_directory(remote_host: str) -> str:
  80. """Resolve the remote user's home directory over SSH."""
  81. home_result = subprocess.run(
  82. ["ssh", remote_host, "printf '%s' \"$HOME\""],
  83. check=False,
  84. capture_output=True,
  85. text=True,
  86. )
  87. if home_result.returncode != 0:
  88. error_output = home_result.stderr.strip() or home_result.stdout.strip() or "SSH home resolution failed"
  89. raise RuntimeError(f"Failed to resolve remote home directory on '{remote_host}': {error_output}")
  90. remote_home = home_result.stdout.strip()
  91. if not remote_home:
  92. raise RuntimeError(f"Failed to resolve remote home directory on '{remote_host}': empty response")
  93. return remote_home
  94. def resolve_remote_absolute_path(remote_host: str, remote_path: str, trailing_slash: bool = False) -> str:
  95. """Resolve ~-prefixed remote paths to absolute remote filesystem paths."""
  96. normalized = remote_path.rstrip("/")
  97. suffix = "/" if trailing_slash else ""
  98. if normalized not in {"~", ""} and not normalized.startswith("~/"):
  99. return f"{normalized}{suffix}"
  100. remote_home = resolve_remote_home_directory(remote_host)
  101. return f"{remote_home}{suffix}" if normalized in {"~", ""} else f"{remote_home}/{normalized[2:]}{suffix}"
  102. def resolve_remote_upload_target(remote_host: str, remote_path: str, trailing_slash: bool = False) -> str:
  103. """Resolve remote paths to absolute scp targets."""
  104. absolute_path = resolve_remote_absolute_path(remote_host, remote_path, trailing_slash=trailing_slash)
  105. return build_scp_remote_target(remote_host, absolute_path, trailing_slash=trailing_slash)
  106. def _write_staging_files(staging_dir: Path, rendered_files: dict[str, str]) -> None:
  107. """Write rendered files to a local staging directory."""
  108. staging_dir.mkdir(parents=True, exist_ok=True)
  109. for file_path, content in rendered_files.items():
  110. full_path = staging_dir / file_path
  111. full_path.parent.mkdir(parents=True, exist_ok=True)
  112. full_path.write_text(content, encoding="utf-8")
  113. def write_rendered_files_remote(
  114. remote_host: str,
  115. remote_path: str,
  116. rendered_files: dict[str, str],
  117. ) -> None:
  118. """Upload rendered files to a remote host over SSH."""
  119. with tempfile.TemporaryDirectory(prefix="boilerplates-remote-") as staging_root:
  120. staging_dir = Path(staging_root)
  121. _write_staging_files(staging_dir, rendered_files)
  122. remote_mkdir_path = build_remote_shell_path(remote_path)
  123. remote_copy_target = resolve_remote_upload_target(remote_host, remote_path, trailing_slash=True)
  124. mkdir_result = subprocess.run(
  125. ["ssh", remote_host, f"mkdir -p -- {remote_mkdir_path}"],
  126. check=False,
  127. capture_output=True,
  128. text=True,
  129. )
  130. if mkdir_result.returncode != 0:
  131. error_output = mkdir_result.stderr.strip() or mkdir_result.stdout.strip() or "SSH mkdir failed"
  132. raise RuntimeError(f"Failed to prepare remote directory '{remote_path}' on '{remote_host}': {error_output}")
  133. upload_result = subprocess.run(
  134. ["scp", "-r", f"{staging_dir}/.", remote_copy_target],
  135. check=False,
  136. capture_output=True,
  137. text=True,
  138. )
  139. if upload_result.returncode != 0:
  140. error_output = upload_result.stderr.strip() or upload_result.stdout.strip() or "SCP upload failed"
  141. raise RuntimeError(f"Failed to upload files to '{remote_host}:{remote_path}': {error_output}")
  142. __all__ = [
  143. "GenerationDestination",
  144. "build_remote_shell_path",
  145. "build_scp_remote_target",
  146. "format_remote_destination",
  147. "normalize_output_path",
  148. "prompt_generation_destination",
  149. "resolve_cli_destination",
  150. "resolve_remote_absolute_path",
  151. "resolve_remote_home_directory",
  152. "resolve_remote_upload_target",
  153. "write_rendered_files_remote",
  154. ]