From b68c042f3292d5ea83dd85a77c818510f3df7076 Mon Sep 17 00:00:00 2001
From: Bingsu <ks2515@naver.com>
Date: Sat, 11 Feb 2023 00:14:22 +0900
Subject: [PATCH] feat: add ssh tunnels

---
 .gitignore            |  3 ++
 README.md             | 10 ++++++-
 preload.py            | 12 ++++++++
 scripts/_util.py      | 11 +++++++
 scripts/ssh_tunnel.py | 68 +++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 103 insertions(+), 1 deletion(-)
 create mode 100644 scripts/_util.py
 create mode 100644 scripts/ssh_tunnel.py

diff --git a/.gitignore b/.gitignore
index 8a434c5..8dd029d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -171,3 +171,6 @@ poetry.toml
 .ruff_cache/
 
 # End of https://www.toptal.com/developers/gitignore/api/python
+
+id_rsa
+id_rsa.pub
diff --git a/README.md b/README.md
index e97e8b0..1970397 100644
--- a/README.md
+++ b/README.md
@@ -4,6 +4,14 @@ Tunneling extension for [AUTOMATIC1111/stable-diffusion-webui](https://github.co
 
 ## Usage
 
-### cloudflared
+### [cloudflared](https://try.cloudflare.com/)
 
 add `--cloudflared` to commandline options.
+
+### [localhost.run](https://localhost.run/)
+
+add `--localhostrun` to commandline options.
+
+### [remote.moe](https://github.com/fasmide/remotemoe)
+
+add `--remotemoe` to commandline options.
diff --git a/preload.py b/preload.py
index 94e3c49..10d066e 100644
--- a/preload.py
+++ b/preload.py
@@ -7,3 +7,15 @@ def preload(parser: argparse.ArgumentParser):
         action="store_true",
         help="use trycloudflare, alternative to gradio --share",
     )
+
+    parser.add_argument(
+        "--localhostrun",
+        action="store_true",
+        help="use localhost.run, alternative to gradio --share",
+    )
+
+    parser.add_argument(
+        "--remotemoe",
+        action="store_true",
+        help="use remote.moe, alternative to gradio --share",
+    )
diff --git a/scripts/_util.py b/scripts/_util.py
new file mode 100644
index 0000000..4f7a9d8
--- /dev/null
+++ b/scripts/_util.py
@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+import shlex
+import subprocess
+from pathlib import Path
+
+
+def gen_key(path: str | Path) -> None:
+    arg_string = f'ssh-keygen -t rsa -b 4096 -N "" -f {path}'
+    args = shlex.split(arg_string)
+    subprocess.run(args, check=True)
diff --git a/scripts/ssh_tunnel.py b/scripts/ssh_tunnel.py
new file mode 100644
index 0000000..34481c5
--- /dev/null
+++ b/scripts/ssh_tunnel.py
@@ -0,0 +1,68 @@
+import atexit
+import re
+import shlex
+import subprocess
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+from modules.shared import cmd_opts
+
+from ._util import gen_key
+
+LOCALHOST_RUN = "localhost.run"
+localhostrun_pattern = re.compile(r"(?P<url>https?://\S+\.lhr\.life)")
+remotemoe_pattern = re.compile(r"(?P<url>https?://\S+\.remote\.moe)")
+
+
+def ssh_tunnel(host: str = LOCALHOST_RUN) -> None:
+    ssh_name = "id_rsa"
+    ssh_path = Path(__file__).parent.parent / ssh_name
+
+    tmp = None
+    if not ssh_path.exists():
+        try:
+            gen_key(ssh_path)
+        # write permission error or etc
+        except subprocess.CalledProcessError:
+            tmp = TemporaryDirectory()
+            ssh_path = Path(tmp.name) / ssh_name
+            gen_key(ssh_path)
+
+    port = cmd_opts.port if cmd_opts.port else 7860
+    arg_string = (
+        f"ssh -R 80:localhost:{port} -o StrictHostKeyChecking=no -i {ssh_path} {host}"
+    )
+    args = shlex.split(arg_string)
+    tunnel = subprocess.Popen(
+        args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, encoding="utf-8"
+    )
+
+    atexit.register(tunnel.terminate)
+    if tmp is not None:
+        atexit.register(tmp.cleanup)
+
+    tunnel_url = ""
+    lines = 27 if host == LOCALHOST_RUN else 5
+    pattern = localhostrun_pattern if host == LOCALHOST_RUN else remotemoe_pattern
+    for _ in range(lines):
+        line = tunnel.stdout.readline()
+        if line.startswith("Warning"):
+            print(line, end="")
+
+        url_match = pattern.search(line)
+        if url_match:
+            tunnel_url = url_match.group("url")
+            break
+    else:
+        raise RuntimeError(f"Failed to run {host}")
+
+    print(f" * Running on {tunnel_url}")
+
+
+if cmd_opts.localhostrun:
+    print("localhost.run detected, trying to connect...")
+    ssh_tunnel(LOCALHOST_RUN)
+
+if cmd_opts.remotemoe:
+    print("remote.moe detected, trying to connect...")
+    ssh_tunnel("remote.moe")