python 又一个Mac SSH SOCKS代理脚本。这个正确处理连接失败。

Posted 2021-05-09

#!/usr/bin/env python
# Public domain.
"""A script to establish an SSH SOCKS proxy and enable it in the OS X network configuration.
Cleanly handles failure to connect. Kills ssh and removes network configuration on CTRL-C.
First argument gives the hostname for use with ssh."""

import subprocess, sys, signal

proxy_command_base = ['networksetup', '-setsocksfirewallproxystate', 'Wi-Fi']

host = sys.argv[1]

ssh = subprocess.Popen([
    'ssh', '-ND', '1080',

    # Have ssh print 'connected' after successful connection
    '-o', 'PermitLocalCommand=yes',
    '-o', "LocalCommand=echo connected",

    host],

    stdout=subprocess.PIPE
    )

try:
    # Only proceed with SOCKS configuration once ssh has connected successfully.
    # If there was an error, it will have gone to stderr and stdout will be empty.
    if ssh.stdout.readline() == "connected\n":

        # Open and keep open a shell with root priviledges. sudo will prompt for password here.
        admin_shell = subprocess.Popen(
                ['sudo', '-s'],

                # Line buffered so we're sure the shell gets our commands immediately.
                bufsize=1,

                stdin=subprocess.PIPE,

                # Don't let SIGINT (ctrl-c) kill the shell - we need to use it to clean up.
                preexec_fn = lambda: signal.signal(signal.SIGINT, signal.SIG_IGN)
                )

        try:
            admin_shell.stdin.write(' '.join(proxy_command_base + ['on']) + \
                    " && echo 'SOCKS proxy via %s on. CTRL-C to turn off.'\n" % host)

            # Block here until ssh exits from SIGINT (ctrl-c) or loss of network connection, etc)
            ssh.wait()

        # SIGINT will have killed ssh; carry on to clean up SOCKS config.
        except KeyboardInterrupt:
            print

        admin_shell.stdin.write(' '.join(proxy_command_base + ['off']) + ' && echo "SOCKS proxy off."\n')
        admin_shell.stdin.close()
        admin_shell.wait()

finally:
    # If something goes wrong with the proxy config, make sure we kill ssh.
    if ssh.poll() is None:
        ssh.kill()