forked from mincomk/rsh
110 lines
3.1 KiB
Rust
110 lines
3.1 KiB
Rust
use axum::extract::{Path, Query};
|
|
use axum::http::{header, HeaderMap, StatusCode};
|
|
use axum::response::{IntoResponse, Response};
|
|
use serde::Deserialize;
|
|
|
|
static STUB_X86_64: &[u8] = include_bytes!(env!("STUB_X86_64_PATH"));
|
|
static STUB_AARCH64: &[u8] = include_bytes!(env!("STUB_AARCH64_PATH"));
|
|
|
|
const SCRIPT_TEMPLATE: &str = r#"#!/bin/sh
|
|
set -eu
|
|
SESSION='__SESSION__'
|
|
BASE='__BASE__'
|
|
WS='__WS__'
|
|
A=$(uname -m)
|
|
case "$A" in
|
|
x86_64|amd64) ARCH=x86_64 ;;
|
|
aarch64|arm64) ARCH=aarch64 ;;
|
|
*) echo "rsh: unsupported arch $A" >&2; exit 1 ;;
|
|
esac
|
|
TMP=$(mktemp /tmp/rsh.XXXXXX)
|
|
trap 'rm -f "$TMP"' EXIT
|
|
if command -v curl >/dev/null 2>&1; then
|
|
curl -fsSL "$BASE/rsh/$ARCH" -o "$TMP"
|
|
elif command -v wget >/dev/null 2>&1; then
|
|
wget -qO "$TMP" "$BASE/rsh/$ARCH"
|
|
else
|
|
echo "rsh: need curl or wget" >&2; exit 1
|
|
fi
|
|
chmod +x "$TMP"
|
|
printf 'Password (empty for none): ' >&2
|
|
stty -echo 2>/dev/null || true
|
|
IFS= read -r PW || PW=''
|
|
stty echo 2>/dev/null || true
|
|
echo >&2
|
|
trap - EXIT
|
|
if [ -n "$PW" ]; then
|
|
exec "$TMP" --url "$WS/ws/stub" --session "$SESSION" --password "$PW"
|
|
else
|
|
exec "$TMP" --url "$WS/ws/stub" --session "$SESSION"
|
|
fi
|
|
"#;
|
|
|
|
#[derive(Deserialize)]
|
|
pub struct RunQuery {
|
|
s: Option<String>,
|
|
}
|
|
|
|
pub async fn run_sh(headers: HeaderMap, Query(q): Query<RunQuery>) -> Response {
|
|
let session = q.s.unwrap_or_else(|| "default".into());
|
|
if !is_valid_session(&session) {
|
|
return (StatusCode::BAD_REQUEST, "invalid session id").into_response();
|
|
}
|
|
let (scheme, ws_scheme) = detect_scheme(&headers);
|
|
let host = headers
|
|
.get(header::HOST)
|
|
.and_then(|v| v.to_str().ok())
|
|
.unwrap_or("localhost");
|
|
let base = format!("{scheme}://{host}");
|
|
let ws = format!("{ws_scheme}://{host}");
|
|
let script = SCRIPT_TEMPLATE
|
|
.replace("__SESSION__", &session)
|
|
.replace("__BASE__", &base)
|
|
.replace("__WS__", &ws);
|
|
(
|
|
[(header::CONTENT_TYPE, "text/x-shellscript; charset=utf-8")],
|
|
script,
|
|
)
|
|
.into_response()
|
|
}
|
|
|
|
pub async fn stub(Path(arch): Path<String>) -> Response {
|
|
let bytes: &[u8] = match arch.as_str() {
|
|
"x86_64" | "amd64" => STUB_X86_64,
|
|
"aarch64" | "arm64" => STUB_AARCH64,
|
|
_ => return StatusCode::NOT_FOUND.into_response(),
|
|
};
|
|
if bytes.is_empty() {
|
|
return (StatusCode::SERVICE_UNAVAILABLE, "stub not embedded in this build").into_response();
|
|
}
|
|
(
|
|
[
|
|
(header::CONTENT_TYPE, "application/octet-stream"),
|
|
(
|
|
header::CONTENT_DISPOSITION,
|
|
"attachment; filename=\"rsh\"",
|
|
),
|
|
],
|
|
bytes,
|
|
)
|
|
.into_response()
|
|
}
|
|
|
|
fn is_valid_session(s: &str) -> bool {
|
|
!s.is_empty()
|
|
&& s.len() <= 64
|
|
&& s.chars().all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-')
|
|
}
|
|
|
|
fn detect_scheme(headers: &HeaderMap) -> (&'static str, &'static str) {
|
|
let proto = headers
|
|
.get("x-forwarded-proto")
|
|
.and_then(|v| v.to_str().ok())
|
|
.unwrap_or("http");
|
|
if proto == "https" {
|
|
("https", "wss")
|
|
} else {
|
|
("http", "ws")
|
|
}
|
|
}
|