proxy.py

 1#!/usr/bin/env python3
 2import os
 3import sys
 4from subprocess import check_output, list2cmdline
 5
 6cwd = os.path.dirname(__file__)
 7ssh_agent = os.path.join(cwd, "ssh-agent.py")
 8proxy_env = "SHELLPROXY_URL"
 9no_proxy_env = "SHELLPROXY_NO_PROXY"
10proxy_config = os.environ.get("SHELLPROXY_CONFIG") or os.path.expandvars("$HOME/.config/proxy")
11
12usage="""shell-proxy: no proxy configuration found.
13
14Set `{env}` or create a config file at `{config}`
15See the plugin README for more information.""".format(env=proxy_env, config=proxy_config)
16
17def get_http_proxy():
18    default_proxy = os.environ.get(proxy_env)
19    no_proxy = os.environ.get(no_proxy_env)
20    if default_proxy and no_proxy:
21        return default_proxy, no_proxy
22
23    if os.path.isfile(proxy_config):
24        proxy_configdata = [line.strip() for line in check_output(proxy_config).decode("utf-8").splitlines()]
25        if len(proxy_configdata) >= 1:
26            if not default_proxy:
27                default_proxy = proxy_configdata[0]
28            if len(proxy_configdata) == 2 and not no_proxy:
29                no_proxy = proxy_configdata[1]
30    
31    if default_proxy:
32        return default_proxy, no_proxy
33    print(usage, file=sys.stderr)
34    sys.exit(1)
35
36
37def make_proxies(url: str, no_proxy: str):
38    proxies = {"%s_PROXY" % _: url for _ in ("HTTP", "HTTPS", "FTP", "RSYNC", "ALL")}
39    proxies.update({name.lower(): value for (name, value) in proxies.items()})
40    proxies["GIT_SSH"] = ssh_agent
41    if no_proxy:
42        proxies.update({"NO_PROXY": no_proxy, "no_proxy": no_proxy})
43    return proxies
44
45
46def merge(mapping: dict):
47    return ("%s=%s" % _ for _ in mapping.items())
48
49
50class CommandSet:
51    proxies = make_proxies(*get_http_proxy())
52    aliases = {
53        _: "env __SSH_PROGRAM_NAME__=%s %s" % (_, ssh_agent)
54        for _ in ("ssh", "sftp", "scp", "slogin", "ssh-copy-id")
55    }
56
57    def enable(self):
58        cmdline("export", *merge(self.proxies))
59        cmdline("alias", *merge(self.aliases))
60
61    def disable(self):
62        cmdline("unset", *self.proxies.keys())
63        cmdline("unalias", *self.aliases.keys())
64
65    def status(self):
66        proxies = (
67            "%11s = %s" % (name, os.environ[name])
68            for name in self.proxies.keys()
69            if name in os.environ
70        )
71        for _ in proxies:
72            cmdline("echo", _)
73
74    def usage(self):
75        print("usage: proxy {enable,disable,status}", file=sys.stderr)
76
77
78def cmdline(*items):
79    print(list2cmdline(items))
80
81
82def main():
83    command = CommandSet()
84    if len(sys.argv) == 1:
85        command.usage()
86        sys.exit(1)
87    getattr(command, sys.argv[1], command.usage)()
88
89
90if __name__ == "__main__":
91    main()