def test_validate(self):
v = Vapid("/tmp/private")
token = "foobar"
sig = v.validate(token)
ok_(v.public_key.verify(
base64.urlsafe_b64decode(sig + '===='[len(sig) % 4:]),
token,
hashfunc=hashlib.sha256))
eq_(v.verify_token(sig, token), True)
def __init__(self,
load_runner,
websocket_url,
statsd_client,
scenario,
endpoint=None,
endpoint_ssl_cert=None,
endpoint_ssl_key=None,
*scenario_args,
**scenario_kw):
self._factory = WebSocketClientFactory(
websocket_url,
headers={"Origin": "http://localhost:9000"})
self._factory.protocol = WSClientProtocol
self._factory.harness = self
if websocket_url.startswith("wss"):
self._factory_context = ssl.ClientContextFactory()
else:
self._factory_context = None
# somewhat bogus encryption headers
self._crypto_key = "keyid=p256dh;dh=c2VuZGVy"
self._encryption = "keyid=p256dh;salt=XZwpw6o37R-6qoZjw6KwAw"
# Processor and Websocket client vars
self._scenario = scenario
self._scenario_args = scenario_args
self._scenario_kw = scenario_kw
self._processors = 0
self._ws_clients = {}
self._connect_waiters = deque()
self._load_runner = load_runner
self._stat_client = statsd_client
self._vapid = Vapid()
if "vapid_private_key" in self._scenario_kw:
self._vapid = Vapid(
private_key=self._scenario_kw.get(
"vapid_private_key"))
else:
self._vapid.generate_keys()
self._claims = None
if "vapid_claims" in self._scenario_kw:
self._claims = self._scenario_kw.get("vapid_claims")
self._endpoint = urlparse.urlparse(endpoint) if endpoint else None
self._agent = None
if endpoint_ssl_cert:
self._agent = Agent(
reactor,
contextFactory=UnverifiedHTTPS(
endpoint_ssl_cert,
endpoint_ssl_key))
def test_sign(self):
v = Vapid(private_key=T_PRIVATE)
claims = {"aud": "example.com", "sub": "*****@*****.**"}
result = v.sign(claims, "id=previous")
eq_(result['Crypto-Key'],
'id=previous,'
'p256ecdsa=' + T_PUBLIC_RAW.strip('='))
items = jws.verify(result['Authorization'][7:],
v.public_key,
algorithms=["ES256"])
eq_(json.loads(items), claims)
result = v.sign(claims)
eq_(result['Crypto-Key'],
'p256ecdsa=' + T_PUBLIC_RAW.strip('='))
def test_save_public_key(self):
v = Vapid()
v.generate_keys()
v.save_public_key("/tmp/p2")
os.unlink("/tmp/p2")
def test_save_private_key(self):
v = Vapid()
v.save_private_key("/tmp/p2")
os.unlink("/tmp/p2")
def test_gen_key(self):
v = Vapid()
v.generate_keys()
ok_(v.public_key)
ok_(v.private_key)
class RunnerHarness(object):
"""Runs multiple instances of a single scenario
Running an instance of the scenario is triggered with :meth:`run`. It
will run to completion or possibly forever.
"""
def __init__(self,
load_runner,
websocket_url,
statsd_client,
scenario,
endpoint=None,
endpoint_ssl_cert=None,
endpoint_ssl_key=None,
*scenario_args,
**scenario_kw):
self._factory = WebSocketClientFactory(
websocket_url,
headers={"Origin": "http://localhost:9000"})
self._factory.protocol = WSClientProtocol
self._factory.harness = self
if websocket_url.startswith("wss"):
self._factory_context = ssl.ClientContextFactory()
else:
self._factory_context = None
# somewhat bogus encryption headers
self._crypto_key = "keyid=p256dh;dh=c2VuZGVy"
self._encryption = "keyid=p256dh;salt=XZwpw6o37R-6qoZjw6KwAw"
# Processor and Websocket client vars
self._scenario = scenario
self._scenario_args = scenario_args
self._scenario_kw = scenario_kw
self._processors = 0
self._ws_clients = {}
self._connect_waiters = deque()
self._load_runner = load_runner
self._stat_client = statsd_client
self._vapid = Vapid()
if "vapid_private_key" in self._scenario_kw:
self._vapid = Vapid(
private_key=self._scenario_kw.get(
"vapid_private_key"))
else:
self._vapid.generate_keys()
self._claims = None
if "vapid_claims" in self._scenario_kw:
self._claims = self._scenario_kw.get("vapid_claims")
self._endpoint = urlparse.urlparse(endpoint) if endpoint else None
self._agent = None
if endpoint_ssl_cert:
self._agent = Agent(
reactor,
contextFactory=UnverifiedHTTPS(
endpoint_ssl_cert,
endpoint_ssl_key))
def run(self):
"""Start registered scenario"""
# Create the processor and start it
processor = CommandProcessor(self._scenario, self._scenario_args, self)
processor.run()
self._processors += 1
def spawn(self, test_plan):
"""Spawn a new test plan"""
self._load_runner.spawn(test_plan)
def connect(self, processor):
"""Start a connection for a processor and queue it for when the
connection is available"""
self._connect_waiters.append(processor)
connectWS(self._factory, contextFactory=self._factory_context)
def send_notification(self, processor, url, data, ttl, claims=None):
"""Send out a notification to a url for a processor"""
url = url.encode("utf-8")
headers = {"TTL": str(ttl)}
crypto_key = self._crypto_key
claims = claims or self._claims
if self._vapid and claims:
headers.update(self._vapid.sign(claims))
crypto_key = "{};p256ecdsa={}".format(
crypto_key,
self._vapid.public_key_urlsafe_base64
)
if data:
headers.update({
"Content-Type": "application/octet-stream",
"Content-Encoding": "aesgcm",
"Crypto-key": crypto_key,
"Encryption": self._encryption,
})
d = treq.post(url,
data=data,
headers=headers,
allow_redirects=False,
agent=self._agent)
d.addCallback(self._sent_notification, processor)
d.addErrback(self._error_notif, processor)
def _sent_notification(self, result, processor):
d = result.content()
d.addCallback(self._finished_notification, result, processor)
d.addErrback(self._error_notif, result, processor)
def _finished_notification(self, result, response, processor):
# Give the fully read content and response to the processor
processor._send_command_result((response, result))
def _error_notif(self, failure, processor):
# Send the failure back
processor._send_command_result((None, failure))
def add_client(self, ws_client):
"""Register a new websocket connection and return a waiting
processor"""
try:
processor = self._connect_waiters.popleft()
except IndexError:
log.msg("No waiting processors for new client connection.")
ws_client.sendClose()
else:
self._ws_clients[ws_client] = processor
return processor
def remove_client(self, ws_client):
"""Remove a websocket connection from the client registry"""
processor = self._ws_clients.pop(ws_client, None)
if not processor:
# Possible failed connection, if we have waiting processors still
# then try a new connection
if len(self._connect_waiters):
connectWS(self._factory, contextFactory=self._factory_context)
return
def remove_processor(self):
"""Remove a completed processor"""
self._processors -= 1
def timer(self, name, duration):
"""Record a metric timer if we have a statsd client"""
self._stat_client.timing(name, duration)
def counter(self, name, count=1):
"""Record a counter if we have a statsd client"""
self._stat_client.increment(name, count)
