websockify/websockify/token_plugins.py

79 lines
2.1 KiB
Python

import os
class BasePlugin(object):
def __init__(self, src):
self.source = src
def lookup(self, token):
return None
class ReadOnlyTokenFile(BasePlugin):
# source is a token file with lines like
# token: host:port
# or a directory of such files
def _load_targets(self):
if os.path.isdir(self.source):
cfg_files = [os.path.join(self.source, f) for
f in os.listdir(self.source)]
else:
cfg_files = [self.source]
self._targets = {}
for f in cfg_files:
for line in [l.strip() for l in open(f).readlines()]:
if line and not line.startswith('#'):
tok, target = line.split(': ')
self._targets[tok] = target.strip().split(':')
def lookup(self, token):
if self._targets is None:
self._load_targets()
if token in self._targets:
return self._targets[token]
else:
return None
# the above one is probably more efficient, but this one is
# more backwards compatible (although in most cases
# ReadOnlyTokenFile should suffice)
class TokenFile(ReadOnlyTokenFile):
# source is a token file with lines like
# token: host:port
# or a directory of such files
def lookup(self, token):
self._load_targets()
return super(TokenFile, self).lookup(token)
class BaseTokenAPI(BasePlugin):
# source is a url with a '%s' in it where the token
# should go
# we import things on demand so that other plugins
# in this file can be used w/o unecessary dependencies
def process_result(self, resp):
return resp.text.split(':')
def lookup(self, token):
import requests
resp = requests.get(self.source % token)
if resp.ok:
return self.process_result(resp)
else:
return None
class JSONTokenApi(BaseTokenAPI):
# source is a url with a '%s' in it where the token
# should go
def process_result(self, resp):
return (resp.json['host'], resp.json['port'])