stow/oh-my-zsh/.oh-my-zsh/plugins/shell-proxy/proxy.py

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