From e75ead7d155d4676d2243b0a05714ad5dee4aab2 Mon Sep 17 00:00:00 2001
From: lamhoangtung <lamhoangtung.vz@gmail.com>
Date: Fri, 25 Oct 2019 02:05:45 +0700
Subject: [PATCH] Add support for SSH publickey authentication

---
 README.md    |  1 +
 remocolab.py | 24 +++++++++++++++++++++---
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 7657c6d..e04c289 100644
--- a/README.md
+++ b/README.md
@@ -21,6 +21,7 @@ It secures TurboVNC connection using SSH port forwarding.
 import remocolab
 remocolab.setupSSHD()
 ```
+Note: Use `remocolab.setupSSHD(use_ssh_key=True)` if you want the server to auto login with your public SSH keys (so you won't need specify which private key to go with which server). In this case, prepare a `authorized_keys` file like [this](https://gist.github.com/lamhoangtung/4fca574da11ef45869bdfea8062417b5/raw/ebdc8c6f8fca2162ae3665f91271bd1fc0fa99b6/authorized_keys)
 
 - SSH and TurboVNC:
 ```python3
diff --git a/remocolab.py b/remocolab.py
index 85c5d89..a287b19 100644
--- a/remocolab.py
+++ b/remocolab.py
@@ -32,7 +32,8 @@ def _check_gpu_available():
 
   return IPython.utils.io.ask_yes_no("Do you want to continue? [y/n]")
 
-def _setupSSHDImpl(ngrok_token, ngrok_region):
+
+def _setupSSHDImpl(ngrok_token, ngrok_region, public_ssh_key):
   #apt-get update
   #apt-get upgrade
   cache = apt.Cache()
@@ -56,6 +57,10 @@ def _setupSSHDImpl(ngrok_token, ngrok_region):
   #Prevent ssh session disconnection.
   with open("/etc/ssh/sshd_config", "a") as f:
     f.write("\n\nClientAliveInterval 120\n")
+    if public_ssh_key is not None:
+      f.write("""AuthorizedKeysFile / root/.ssh/authorized_keys" >> /etc/ssh/sshd_config\n""")
+      f.write(""""PubkeyAuthentication yes" >> /etc/ssh/sshd_config\n""")
+
 
   print("ECDSA key fingerprint of host:")
   ret = subprocess.run(
@@ -79,6 +84,10 @@ def _setupSSHDImpl(ngrok_token, ngrok_region):
   subprocess.run(["useradd", "-s", "/bin/bash", "-m", user_name])
   subprocess.run(["chpasswd"], input = f"root:{root_password}", universal_newlines = True)
   subprocess.run(["chpasswd"], input = f"{user_name}:{user_password}", universal_newlines = True)
+  if public_ssh_key is not None:
+    subprocess.run(["wget", "-P", "/root/.ssh/", public_ssh_key])
+    subprocess.run(["chmod", "700" , "/root/.ssh"])
+    subprocess.run(["chmod", "600", "/root/.ssh/authorized_keys"])
   subprocess.run(["service", "ssh", "restart"])
 
   if not pathlib.Path('/root/.ngrok2/ngrok.yml').exists():
@@ -106,10 +115,19 @@ def _setupSSHDImpl(ngrok_token, ngrok_region):
   print(f"ssh {ssh_common_options} -L 5901:localhost:5901 -p {port} {user_name}@{hostname}")
   print("✂️"*24)
 
-def setupSSHD(ngrok_region = None, check_gpu_available = False):
+
+def setupSSHD(ngrok_region=None, check_gpu_available=False, use_ssh_key=False):
   if check_gpu_available and not _check_gpu_available():
     return False
 
+  if use_ssh_key:
+    print("---")
+    print("Copy&paste link to your raw authorized public SSH key")
+    print("Example: https://gist.github.com/lamhoangtung/4fca574da11ef45869bdfea8062417b5/raw/ebdc8c6f8fca2162ae3665f91271bd1fc0fa99b6/authorized_keys")
+    public_ssh_key = getpass.getpass()
+  else:
+    public_ssh_key = None
+
   print("---")
   print("Copy&paste your tunnel authtoken from https://dashboard.ngrok.com/auth")
   print("(You need to sign up for ngrok and login,)")
@@ -127,7 +145,7 @@ def setupSSHD(ngrok_region = None, check_gpu_available = False):
     print("in - India (Mumbai)")
     ngrok_region = region = input()
 
-  _setupSSHDImpl(ngrok_token, ngrok_region)
+  _setupSSHDImpl(ngrok_token, ngrok_region, public_ssh_key)
   return True
 
 def _setupVNC():