chrome_ssh.sh

#!/usr/bin/env python3.7
import argparse
import os
import shlex
import signal
import socket
import sys
from contextlib import closing
from subprocess import Popen, CalledProcessError
from tempfile import TemporaryDirectory
from time import sleep

if __name__ != "__main__":
    raise Exception(
        "This file is intended to be executed as an executable program. You cannot use it as a module."
        "To run this script, run the ./build command"
    )


def get_chrome_path():
    mac_os_path = '/Applications/Google Chrome.app/Contents/MacOS/Google Chrome'
    if os.path.exists(mac_os_path):
        return mac_os_path
    system_path = os.environ.get("PATH", "").split(os.pathsep)

    possible_path = (
        os.path.join(path_elem, "google-chrome")
        for path_elem in system_path
    )
    existing_path = (
        os.path.exists(path_elem) for path_elem in possible_path
    )
    try:
        return next(existing_path)
    except StopIteration:
        raise Exception(
            "You need to install Google Chrome to use this functionality."
            ""
            "For more information look at:"
            "https://support.google.com/chrome/answer/95346"
        )


def shlex_cmd(cmd):
    return " ".join([shlex.quote(c) for c in cmd])


def find_free_port():
    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
        s.bind(('', 0))
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        return s.getsockname()[1]


def start_chrome_with_ssh_forward(gcloud_base, gcloud_extra):
    proxy_port = find_free_port()
    print("Found free port:", proxy_port)

    ssh_proc = None
    chrome_proc = None

    def kill_subprocesses(int, frame):
        if ssh_proc:
            ssh_proc.kill()
        if chrome_proc:
            chrome_proc.kill()

    signal.signal(signal.SIGINT, kill_subprocesses)
    signal.signal(signal.SIGTERM, kill_subprocesses)
    with TemporaryDirectory() as chrome_profile_tmp:
        ssh_cmd = [*gcloud_base, 'ssh', f'--ssh-flag=-D {proxy_port}', '--ssh-flag=-N']
        if gcloud_extra:
            ssh_cmd.extend(gcloud_extra)
        print("Executing command: ", shlex_cmd(ssh_cmd), file=sys.stderr)
        ssh_proc = Popen(ssh_cmd)
        sleep(5)
        if ssh_proc.poll() is not None:
            print(ssh_proc.returncode, ssh_proc.args, ssh_proc.stdout, ssh_proc.stderr)
            raise CalledProcessError(
                ssh_proc.returncode, ssh_proc.args, None, None
            )

        chrome_cmd = [
            get_chrome_path(),
            f'--proxy-server=socks5://localhost:{proxy_port}',
            f'--user-data-dir={chrome_profile_tmp}',
            '--no-default-browser-check',
            '--no-first-run',
            'https://httpbin.org/get'
        ]
        print("Executing command: ", shlex_cmd(chrome_cmd), file=sys.stderr)
        chrome_proc = Popen(chrome_cmd)
        chrome_proc.wait()
        ssh_proc.kill()


def ger_parser():
    parser = argparse.ArgumentParser(prog='chrome_sssh')
    parser.add_argument(dest='cmd')
    parser.add_argument('extra_args', nargs=argparse.REMAINDER)
    return parser


parser = ger_parser()
args = parser.parse_args()

if '--help' in args.extra_args or '-h' in args.extra_args:
    parser.print_help()
    sys.exit(1)

if args.cmd == 'cloud-shell':
    command_base = ['gcloud', 'alpha', 'cloud-shell']
    start_chrome_with_ssh_forward(command_base, args.extra_args)
elif args.cmd == 'compute':
    command_base = ['gcloud', 'compute']
    start_chrome_with_ssh_forward(command_base, args.extra_args)
else:
    raise Exception("Invalid command")