Merge pull request #7 from lamhoangtung/mattermost

Switch Microsoft Teams webhook to Mattermost webhook

.github/workflows/python-package.yml

@@ -14,8 +14,9 @@ jobs:
 
     runs-on: ubuntu-18.04
     strategy:
+      fail-fast: false
       matrix:
-        python-version: ['3.6', '3.7']
+        python-version: ['3.6', '3.7', '3.8', '3.9']
 
     steps:
     - uses: actions/checkout@v2
@@ -43,7 +44,7 @@ jobs:
         flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
     - name: Test with unnitest
       env:
-        TEAMS_WEBHOOK_ADDRESS: ${{ secrets.TEAMS_WEBHOOK_ADDRESS }}
+        MATTERMOST_WEBHOOK_ADDRESS: ${{ secrets.MATTERMOST_WEBHOOK_ADDRESS }}
         TEST_SSH_PUBLIC_KEY: ${{ secrets.TEST_SSH_PUBLIC_KEY }}
       run: |
         sudo -E python3 -m coverage run --source=. -m unittest discover

README.md

@@ -28,6 +28,8 @@ loop_forever()
 
 You can use list of public key or link to a raw text file of `authorized_keys` like [this](https://gist.githubusercontent.com/lamhoangtung/4fca574da11ef45869bdfea8062417b5/raw/320893c60a5a150f61481899201664761136fae7/authorized_keys) as well
 
+Optinally, you can also specify a Mattermost Webhook URL with `mattermost_webhook_address` when calling `setup_ssh` to send a push notification to your Mattermost channel when the SSH tunel is ready.
+
 Run it, after about 2 minutes, you will see something like this:
 
 ```bash

colab_ssh/__init__.py

@@ -1,3 +1,4 @@
 # pylint: disable=missing-module-docstring
-from .main import setup_ssh, loop_forever
+from .main import loop_forever, setup_ssh
+
 __all__ = ('setup_ssh', 'loop_forever')  # pragma: no cover

colab_ssh/main.py

@@ -1,16 +1,17 @@
 """Just look at the name, it's main"""
 import time
+from typing import List, Optional, Union
 
 from colab_ssh.config import config_root_password, install_common_tool
-from colab_ssh.notification import send_notification_to_microsoft_teams
+from colab_ssh.notification import send_notification_to_mattermost
 from colab_ssh.ssh import config_ssh_server, parse_public_key
 from colab_ssh.tunel import config_argo_tunnel
 from colab_ssh.utils import check_gpu_available, get_instance_info
 
 
-def setup_ssh(public_key, teams_webhook_address: str = None):
+def setup_ssh(public_key: Union[str, List[str]], mattermost_webhook_address: Optional[str] = None):
     """
-    Setup an SSH tunel to the current Colab notebook instance with ssh public key authentication
+    Setup an SSH tunel to the current Colab notebook instance with SSH public key authentication
 
     Parameters:
         public_key:
@@ -19,7 +20,7 @@ def setup_ssh(public_key, teams_webhook_address: str = None):
             - (str): Link to a text file (authorized_keys) that cotains all the public keys that will be
             able to authenticate the SSH connection
         webhook_address:
-            - (str): The webhook address for microsoft teams for push notification
+            - (str): The webhook address of Mattermost for push notification
 
     After about 2 minutes of running, the bash command to initialize the SSH connection will be print out
     """
@@ -39,14 +40,15 @@ def setup_ssh(public_key, teams_webhook_address: str = None):
     install_common_tool()
 
     # Config Argo Tunnel
-    msg, ssh_command, hostname = config_argo_tunnel(msg)
+    msg, ssh_command, ssh_config, hostname = config_argo_tunnel(msg)
 
-    # Send notification to Microsoft Teams
-    if teams_webhook_address is not None:
+    # Send notification to Mattermost
+    if mattermost_webhook_address is not None:
         spec = get_instance_info()
         spec['ssh_command'] = ssh_command
+        spec['ssh_config'] = ssh_config
         spec['hostname'] = hostname
-        send_notification_to_microsoft_teams(teams_webhook_address, spec)
+        send_notification_to_mattermost(mattermost_webhook_address, spec)
 
     print(msg)
 

colab_ssh/notification.py

@@ -1,32 +1,46 @@
-"""Send push notification to Microsoft Teams using webhook. Just to manage all the instance"""
+"""Send push notification to Mattermost Channel using webhook. Just to manage all the instance"""
 
+import json
 import os
 from typing import Dict
 
-import pymsteams
+import requests
 
+COLAB_USER_NAME = 'colab'
+COLAB_USER_ICON_LINK = 'https://colab.research.google.com/img/colab_favicon_256px.png'
 
-def send_notification_to_microsoft_teams(webhook_address: str, spec: Dict):
+
+def send_notification_to_mattermost(webhook_address: str, spec: Dict[str, str]):
     """
-    Send push notification to Microsoft Teams using webhook. Just to manage all the instance
+    Send push notification to Mattermost using webhook. Just to manage all the instance
     """
-
-    message = pymsteams.connectorcard(webhook_address)
+    text = ""
     if os.environ.get("IS_TESTING_CI") is not None:
-        message.text(
-            "Hi @channel, a new CI/CD test for the colab ssh pip package was initialized! Here was it configuration:")
+        text += "Hi @channel, a new CI/CD test for the Colab SSH pip package was initialized! Here was it configuration:"
     else:  # pragma: no cover
-        message.text(
-            "Hi @channel, a new colab spot instance was created! Here was it configuration:")
-
-    section = pymsteams.cardsection()
-    section.addFact("CPU", spec['cpu'])
-    section.addFact("RAM", spec['ram'])
-    section.addFact("GPU", spec['gpu'])
-    section.addFact("Hostname", spec['hostname'])
-    section.addFact("Connection command", spec['ssh_command'])
-    message.addSection(section)
+        text += "Hi @channel, a new Colab spot instance was created :tada::tada::tada:! Here was it configuration:"
+    text += "\n\n"
+    text += f"| **CPU**      | {spec['cpu']}                                                  |\n"
+    text += f"|--------------|----------------------------------------------------------------|\n"
+    text += f"| **RAM**      | {spec['ram']}                                                  |\n"
+    text += f"| **GPU**      | {spec['gpu']}                                                  |\n"
+    text += f"| **Hostname** | `{spec['hostname']}`                                           |\n"
+    text += "\n"
+    text += "To **connect** to it, use the following configuration in your `~/.ssh/config` file:\n"
+    text += f"```ssh-config\n{spec['ssh_config']}\n```\n"
+    text += "Don't forget to **comment** to this post to **claim your colab instance** now, these thing don't really grow on tree ;). ***Happy coding!***\n"
+    payload = {
+        "username": COLAB_USER_NAME,
+        "icon_url": COLAB_USER_ICON_LINK,
+        "text": text
+    }
+    payload = json.dumps(payload)
+    headers = {
+        'Content-Type': 'application/json'
+    }
     try:
-        message.send()
-    except Exception as exception:  # pylint: disable=broad-except 
-        print(f"Error sending notification to Microsoft Teams: {exception}")
+        response = requests.request(
+            "POST", webhook_address, headers=headers, data=payload)
+        response.raise_for_status()
+    except Exception as ex:
+        print(f"Cannot send notification to Mattermost: {ex}")

colab_ssh/progress_bar.py

@@ -6,10 +6,11 @@
 https://apt-team.pages.debian.net/python-apt/library/index.html
 """
 
-from IPython.core.display import display
 import apt
 import apt.debfile
 import ipywidgets
+from IPython.core.display import display
+
 
 class NoteProgress(apt.progress.base.InstallProgress, apt.progress.base.AcquireProgress, apt.progress.base.OpProgress):
     """

colab_ssh/tunel.py

@@ -3,10 +3,12 @@
 import subprocess
 import time
 import urllib
+from typing import Tuple
+
 from colab_ssh.utils import download_file
 
 
-def config_argo_tunnel(msg: str):
+def config_argo_tunnel(msg: str) -> Tuple[str, str, str, str]:
     """
     Config Argo tunnel for the SSH tunnel
 
@@ -15,8 +17,9 @@ def config_argo_tunnel(msg: str):
 
     Return:
         msg (str): The message after added tunnel information and ssh command
-        ssh_command (str): The SSH command for Microsoft Teams push notification
-        hostname (str): The hostname of the server, also for Microsoft Teams push notification
+        ssh_command (str): The SSH command for Mattermost push notification
+        ssh_config (str): The SSH config block for Mattermost push notification
+        hostname (str): The hostname of the server, also for Mattermost push notification
     """
     download_file(
         "https://bin.equinox.io/c/VdrWdbjqyF/cloudflared-stable-linux-amd64.tgz", "cloudflared.tgz")
@@ -51,7 +54,8 @@ def config_argo_tunnel(msg: str):
             hostname = text[begin + len(sub): end]
             break
     if hostname is None:
-        raise RuntimeError("Failed to get user hostname from cloudflared")  # pragma: no cover
+        raise RuntimeError(
+            "Failed to get user hostname from cloudflared")  # pragma: no cover
 
     ssh_common_options = "-o UserKnownHostsFile=/dev/null -o VisualHostKey=yes"
     ssh_common_options += " -oProxyCommand=\"cloudflared access ssh --hostname %h\""
@@ -64,8 +68,9 @@ def config_argo_tunnel(msg: str):
     msg += "✂️"*24 + "\n"
     msg += "Or you can use the following configuration in your .ssh/config file:\n"
     msg += "✂️"*24 + "\n"
-    msg += f"Host colab\n\tHostName {hostname}\n\tUser root\n\tUserKnownHostsFile /dev/null\n"
-    msg += "\tVisualHostKey yes\n\tStrictHostKeyChecking no\n"
-    msg += "\tProxyCommand cloudflared access ssh --hostname %h\n"
+    ssh_config = f"Host colab\n\tHostName {hostname}\n\tUser root\n\tUserKnownHostsFile /dev/null\n"
+    ssh_config += "\tVisualHostKey yes\n\tStrictHostKeyChecking no\n"
+    ssh_config += "\tProxyCommand cloudflared access ssh --hostname %h\n"
+    msg += ssh_config
     msg += "✂️"*24 + "\n"
-    return msg, ssh_command, hostname
+    return msg, ssh_command, ssh_config, hostname

colab_ssh/utils.py

@@ -1,16 +1,16 @@
 """Contain utils function"""
 
-import os
 import multiprocessing
+import os
 import shutil
 import subprocess
 import urllib.request
 from typing import Dict
 
+import apt
 import IPython.utils.io
 from IPython.core.getipython import get_ipython
 from psutil import virtual_memory
-import apt
 
 from colab_ssh.progress_bar import NoteProgress
 
@@ -78,7 +78,7 @@ def run_command(setup_script: str):
             get_ipython().system_raw(command)  # pragma: no cover
 
 
-def get_instance_info() -> Dict:
+def get_instance_info() -> Dict[str, str]:
     """
     Get current instance information
 
@@ -104,18 +104,18 @@ def get_instance_info() -> Dict:
         possible_gpu_type = ['P100', 'K80', 'T4', 'V100']
         for each in possible_gpu_type:
             if each in gpu_info:
-                gpu_type = 'NVIDIA Tesla {}'.format(each)
+                gpu_type = f'NVIDIA **Tesla {each}**'
                 break
         gpu_info = [each for each in raw_gpu_info if 'MiB' in each]
         gpu_vram_gb = int(gpu_info[0].split(
             '|')[-3].split('/')[-1].strip().replace('MiB', ''))/1024
         if gpu_type is not None:
-            gpu_type = "{} - {:.2f} GB".format(gpu_type, gpu_vram_gb)
+            gpu_type = "{} - **{:.2f}** GB".format(gpu_type, gpu_vram_gb)
     total_ram = virtual_memory().total / 1e9
     return {
-        'cpu':  f"{multiprocessing.cpu_count()} cores",
+        'cpu':  f"**{multiprocessing.cpu_count()}** cores",
         'gpu': gpu_type,
-        'ram': "{:.2f} GB".format(total_ram),
+        'ram': "**{:.2f}** GB".format(total_ram),
     }
 
 

setup.py

@@ -1,35 +1,28 @@
 #!/usr/bin/env python
 import os
-from setuptools import setup, find_packages
-__version__ = "0.1.4"
 
+from setuptools import find_packages, setup
 
-def readme():
-    with open(os.path.join(os.path.dirname(__file__), 'README.md')) as f:
-        return f.read()
-
+__version__ = "0.1.5"
 
-def parse_requirements(filename):
-    """ load requirements from a pip requirements file """
-    lineiter = (line.strip() for line in open(filename))
-    return [line for line in lineiter if line and not line.startswith("#")]
 
+def read_me():
+    with open(os.path.join(os.path.dirname(__file__), 'README.md')) as f:
+        return f.read()
 
-install_requires = parse_requirements('requirements.txt')
 
 setup(
     name="linus_colab_ssh",
     version=__version__,
     description='Create SSH tunel to a running colab notebook',
-    long_description=readme(),
+    long_description=read_me(),
     long_description_content_type='text/markdown',
     url='https://github.com/lamhoangtung/colab_ssh',
     author='Hoang Tung Lam',
     author_email='[email protected]',
     include_package_data=True,
     packages=find_packages(
         exclude=["*.tests", "*.tests.*", "tests.*", "tests"]),
-    install_requires=install_requires,
     keywords=['ssh', 'colab'],
     setup_requires=[],
     dependency_links=[],
@@ -41,5 +34,7 @@ def parse_requirements(filename):
         'License :: OSI Approved :: MIT License',
         'Programming Language :: Python :: 3.6',
         'Programming Language :: Python :: 3.7',
+        'Programming Language :: Python :: 3.8',
+        'Programming Language :: Python :: 3.9',
     ],
 )

tests/test_setup.py

@@ -1,12 +1,12 @@
+import os
 import unittest
 from unittest import mock
-import os
 
-from psutil import test
 from colab_ssh import setup_ssh
-from colab_ssh.notification import send_notification_to_microsoft_teams
+from colab_ssh.notification import send_notification_to_mattermost
 from colab_ssh.ssh import parse_public_key
 
+
 @mock.patch.dict(os.environ, {"IS_TESTING_CI": "TRUE"})
 class TestSetup(unittest.TestCase):
     def __init__(self, *args, **kwargs):
@@ -15,7 +15,7 @@ def __init__(self, *args, **kwargs):
 
     def test_setup(self):
         # Test setup SSH with notification
-        webhook_address = os.environ.get("TEAMS_WEBHOOK_ADDRESS")
+        webhook_address = os.environ.get("MATTERMOST_WEBHOOK_ADDRESS")
         ssh_public_key = os.environ.get("TEST_SSH_PUBLIC_KEY")
         if webhook_address is None:
             print("Lol this shit is none")
@@ -27,10 +27,10 @@ def test_fail_push_notification(self):
             'cpu': 'hihi',
             'gpu': 'haha',
             'ram': 'hoho',
-            'ssh_command': 'lul',
+            'ssh_config': 'lul',
             'hostname': 'lol'
         }
-        send_notification_to_microsoft_teams(webhook_address, spec)
+        send_notification_to_mattermost(webhook_address, spec)
 
     def test_parse_public_key(self):
         test_key = ['ssh-a', 'ssh-b', 'ssh-c']
@@ -44,5 +44,6 @@ def test_parse_public_key(self):
         except ValueError as ex:
             pass
 
+
 if __name__ == '__main__':
     unittest.main()