tunnel start/stop function (#6)

Author: Leemoonsoo

requirements.txt

@@ -1,2 +1,3 @@
 pyyaml
 requests
+python3-wget

setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="staroid", # Replace with your own username
-    version="0.0.5",
+    version="0.0.6",
     license='MIT',
     author="Staroid",
     author_email="support@staroid.com",
@@ -41,6 +41,6 @@
     install_requires=[
         'requests',
         'pyyaml',
-
+        'python3-wget'
     ]
 )

staroid/cluster.py

@@ -100,5 +100,5 @@ def delete(self, name):
         if r.status_code == 200:
             return cluster_to_del
         else:
-            logging.error("Can not create clusters {}", r.status_code)
+            logging.error("Can not delete clusters {}", r.status_code)
             return None

staroid/namespace.py

@@ -1,6 +1,8 @@
 import requests
 import logging
 import json
+import subprocess
+import atexit
 from .commit import Commit
 
 class Namespace:
@@ -37,6 +39,7 @@ class NamespaceApi:
     def __init__(self, staroid, cluster):
         self.__staroid = staroid
         self.__cluster = cluster
+        self.__tunnel_processes = {}
 
     def get_all(self):
         r = self.__staroid._api_get(
@@ -181,7 +184,7 @@ def shell_start(self, instance_name):
         )
         if r.status_code == 200:
             js = json.loads(r.text)
-            return Namespace(js)
+            return js
         else:
             logging.error("Can not start shell {}", r.status_code)
             return None
@@ -220,3 +223,64 @@ def get_all_resources(self, instance_name):
         else:
             logging.error("Can not get namespace resources {}", r.status_code)
             return None
+
+    def _is_tunnel_running(self, instance_name):
+        if instance_name in self.__tunnel_processes:
+            p = self.__tunnel_processes[instance_name]
+            p.poll()
+            return p.returncode == None
+        else:
+            return False
+
+    def start_tunnel(self, instance_name, tunnels):
+        if self._is_tunnel_running(instance_name):
+            return
+
+        chisel_path = self.__staroid.get_chisel_path()
+
+        ns = self.get(instance_name)
+        if ns == None:
+            return None
+        resources = self.get_all_resources(instance_name)
+
+        shell_service = None
+        for s in resources["services"]["items"]:
+            if "labels" in s["metadata"]:
+                if "resource.staroid.com/system" in s["metadata"]["labels"]:
+                    if s["metadata"]["labels"]["resource.staroid.com/system"] == "shell":
+                        shell_service = s
+                        break                        
+
+        if shell_service == None:
+            raise Exception("Shell service not found")
+
+        tunnel_server = "https://p{}-{}--{}".format("57682", shell_service["metadata"]["name"], ns.url()[len("https://"):])
+        cmd = [
+            chisel_path,
+            "client",
+            "--header",
+            "Authorization: token {}".format(self.__staroid.get_access_token()),
+            "--keepalive",
+            "10s",
+            tunnel_server
+        ]
+        cmd.extend(tunnels)
+        self.__tunnel_processes[instance_name]=subprocess.Popen(cmd)
+        atexit.register(self.__cleanup)
+
+    def __cleanup(self):
+        timeout_sec = 5
+        for p in self.__tunnel_processes.values(): # list of your processes
+            p_sec = 0
+            for second in range(timeout_sec):
+                if p.poll() == None:
+                    time.sleep(1)
+                    p_sec += 1
+            if p_sec >= timeout_sec:
+                p.kill() # supported from python 2.6
+
+    def stop_tunnel(self, instance_name):
+        if self._is_tunnel_running(instance_name):
+            self.__tunnel_processes[instance_name].kill()
+            del self.__tunnel_processes[instance_name]
+

staroid/staroid.py

@@ -1,11 +1,23 @@
-import os
+import os, stat
 import yaml
 import logging
 import requests
 import json
+from pathlib import Path
+from shutil import which
+import platform
+import subprocess
+import wget
+
 from .cluster import ClusterApi
 from .namespace import NamespaceApi
 
+CHISEL_VERSION="1.6.0"
+CHISEL_ARCH_MAP={
+    "x86_64": "amd64",
+    "i386": "386"
+}
+
 class Org:
     def __init__(self, json):
         self.__json = json
@@ -36,10 +48,15 @@ def principal(self):
 class Staroid:
     """Staroid client object"""
 
-    def __init__(self, access_token=None, account=None, config_path="~/.staroid/config.yaml"):
+    def __init__(self, access_token=None, account=None, config_path=None, cache_dir=None, chisel_path=None):
         self.__api_addr = "https://staroid.com/api"
         self.__access_token = None
         self.__account = None
+        self.__cache_dir = cache_dir
+        self.__chisel_path = chisel_path
+
+        if self.__cache_dir == None:
+            self.__cache_dir = "{}/.staroid".format(str(Path.home()))
 
         # 1. set from configs
         self.__read_config(config_path)
@@ -65,6 +82,9 @@ def __init__(self, access_token=None, account=None, config_path="~/.staroid/conf
                 self.__account = "{}/{}".format(user.provider(), user.principal())
 
     def __read_config(self, config_path):
+        if config_path == None:
+            config_path = "{}/.staroid/config.yaml".format(str(Path.home()))
+
         try:
             with open(config_path, "r") as f:
                 logging.info("Read configuration from " + config_path)
@@ -74,6 +94,52 @@ def __read_config(self, config_path):
         except EnvironmentError:
             pass
 
+    def create_or_get_cache_dir(self, module = ""):
+        "create (if not exists) or return cache dir path for module"
+        cache_dir = "{}/{}".format(self.__cache_dir, module)
+        if not os.path.exists(cache_dir):
+            os.makedirs(cache_dir)
+        return cache_dir
+
+    def __check_cmd(self, cmd):
+        if which(cmd) == None:
+            raise Exception("'{}' command not found".format(cmd))
+
+    def __download_chisel_if_not_exists(self):
+        # check gunzip available
+        self.__check_cmd("gunzip")
+
+        if self.__chisel_path == None:
+            # download chisel binary for secure tunnel if not exists
+            uname = platform.uname()
+            uname.system.lower()
+            if uname.machine not in CHISEL_ARCH_MAP.keys():
+                raise Exception("Can not download chisel automatically. Please download manually from 'https://github.com/jpillora/chisel/releases/download/v{}' and set 'chisel_path' argument".format(CHISEL_VERSION))
+
+            download_url = "https://github.com/jpillora/chisel/releases/download/v{}/chisel_{}_{}_{}.gz".format(
+                CHISEL_VERSION, CHISEL_VERSION, uname.system.lower(), CHISEL_ARCH_MAP[uname.machine])
+            cache_bin = self.create_or_get_cache_dir("bin")
+            chisel_path = "{}/chisel".format(cache_bin)
+
+            if not os.path.exists(chisel_path):
+                # download
+                filename = wget.download(download_url, cache_bin)
+
+                # extract
+                subprocess.run(["gunzip", "-f", filename])
+
+                # rename
+                subprocess.run(["mv", filename.replace(".gz", ""), chisel_path])
+
+                # chmod
+                os.chmod(chisel_path, stat.S_IRWXU)
+
+            self.__chisel_path = chisel_path
+
+    def get_chisel_path(self):
+        self.__download_chisel_if_not_exists()
+        return self.__chisel_path
+
     def cluster(self):
         c = ClusterApi(self)
         return c

tests/test_namespace.py

@@ -37,6 +37,12 @@ def test_crud_namespace(self):
         resources = ns_api.get_all_resources("instance1")
         self.assertTrue(len(resources["services"]) > 0)
 
+        # start tunnel
+        ns_api.start_tunnel("instance1", ["57683:localhost:57683"])
+
+        # stop tunnel
+        ns_api.stop_tunnel("instance1")
+
         # stop shell
         ns_api.shell_stop("instance1")
 

tests/test_staroid.py

@@ -1,6 +1,8 @@
 import unittest
 import tempfile
 import os
+import pathlib
+import shutil
 
 from staroid import Staroid
 
@@ -40,6 +42,22 @@ def test_read_config(self):
         if ac != None:
             os.environ["STAROID_ACCOUNT"] = ac
 
+    def test_download_chisel(self):
+        # given
+        tmp_dir = tempfile.mkdtemp()
+        s = Staroid(cache_dir=tmp_dir)
+
+        # when
+        chisel_path = s.get_chisel_path()
+
+        # then
+        self.assertIsNotNone(chisel_path)
+        self.assertTrue(os.path.isfile(chisel_path))
+        
+        # clean up
+        shutil.rmtree(pathlib.Path(tmp_dir))
+
+
     @unittest.skipUnless(integration_test_ready(), "Integration test environment is not configured")
     def test_read_default_account(self):
         # given access_token is set but account is not set