proxy.py 2.1 KB

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