def test_init(self):
args = TrialSettings()
args.vapid_key = os.path.join(tempfile.gettempdir(), uuid.uuid4().hex)
vapid = Vapid()
vapid.generate_keys()
vapid.save_key(args.vapid_key)
client = PushClient(loop=self.loop, args=args)
assert (client.vapid.public_key.public_numbers().encode_point ==
vapid.public_key.public_numbers().encode_point)
os.unlink(args.vapid_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):
logging.debug("Connecting to {}".format(websocket_url))
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 = ()
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))
if hasattr(endpoint_ssl_cert, 'seek'):
endpoint_ssl_cert.seek(0)
if endpoint_ssl_key and hasattr(endpoint_ssl_key, 'seek'):
endpoint_ssl_key.seek(0)
def run(self):
"""Start registered scenario"""
# Create the processor and start it
processor = CommandProcessor(self._scenario, self._scenario_args,
self._scenario_kw, 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,
headers=None,
claims=None):
"""Send out a notification to a url for a processor
This uses the older `aesgcm` format.
"""
if not headers:
headers = {}
url = url.encode("utf-8")
if "TTL" not in headers:
headers["TTL"] = "0"
crypto_key = self._crypto_key
if claims is None:
claims = ()
claims = claims or self._claims
if self._vapid and claims:
if isinstance(claims, str):
claims = json.loads(claims)
if "aud" not in claims:
# Construct a valid `aud` from the known endpoint
parsed = urlparse.urlparse(url)
claims["aud"] = "{scheme}://{netloc}".format(
scheme=parsed.scheme, netloc=parsed.netloc)
log.msg("Setting VAPID 'aud' to {}".format(claims["aud"]))
headers.update(self._vapid.sign(claims, self._crypto_key))
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)
class PushClient(object):
"""Smoke Test the Autopush push server"""
def __init__(self, args, loop,
tasks: List[List[Union[AnyStr, Dict]]]=None):
self.config = args
self.loop = loop
self.connection = None
self.pushEndpoint = None
self.channelID = None
self.notifications = []
self.uaid = None
self.recv = []
self.tasks = tasks or []
self.output = None
self.vapid_cache = {}
if args.vapid_key:
self.vapid = Vapid().from_file(args.vapid_key)
else:
self.vapid = Vapid()
self.vapid.generate_keys()
self.tls_conn = None
if args.partner_endpoint_cert:
if os.path.isfile(args.partner_endpoint_cert):
context = ssl.create_default_context(
cafile=args.partner_endpoint_cert)
else:
context = ssl.create_default_context(
cadata=args.partner_endpoint_cert)
context.verify_mode = ssl.CERT_REQUIRED
context.check_hostname = False
self.tls_conn = aiohttp.TCPConnector(ssl_context=context)
def _fix_endpoint(self, endpoint: str) -> str:
"""Adjust the endpoint if needed"""
if self.config.partner_endpoint:
orig_path = urlparse(endpoint).path
partner = urlparse(self.config.partner_endpoint)
return "{scheme}://{host}{path}".format(
scheme=partner.scheme,
host=partner.netloc,
path=orig_path)
return endpoint
def _cache_sign(self, claims: Dict[str, str]) -> Dict[str, str]:
"""Pull a VAPID header from the cache or sign the new header
:param claims: list of VAPID claims.
:returns: dictionary of VAPID headers.
"""
vhash = hashlib.sha1()
vhash.update(json.dumps(claims).encode())
key = vhash.digest()
if key not in self.vapid_cache:
self.vapid_cache[key] = self.vapid.sign(claims)
return self.vapid_cache[key]
async def _next_task(self):
"""Tasks are shared between active "cmd_*" commands
and async "recv_*" events. Since both are reading off
the same stack, we centralize that here.
"""
try:
task = self.tasks.pop(0)
logging.debug(">>> cmd_{}".format(task[0]))
await getattr(self, "cmd_" + task[0])(**(task[1]))
return True
except IndexError:
await self.cmd_done()
return False
except PushException:
raise
except AttributeError:
raise PushException("Invalid command: {}".format(task[0]))
except Exception: # pragma nocover
traceback.print_exc()
raise
async def run(self,
server: str="wss://push.services.mozilla.com",
tasks: List[List[Union[AnyStr, Dict]]]=None):
"""Connect to a remote server and execute the tasks
:param server: URL to the Push Server
:param tasks: List of tasks and arguments to run
"""
if tasks:
self.tasks = tasks
if not self.connection:
await self.cmd_connect(server)
while await self._next_task():
pass
async def process(self, message: Dict[str, Any]):
"""Process an incoming websocket message
:param message: JSON message content
"""
mtype = "recv_" + message.get('messageType').lower()
try:
await getattr(self, mtype)(**message)
except AttributeError as ex:
raise PushException(
"Unknown messageType: {}".format(mtype)) from ex
async def receiver(self):
"""Receiver handler for websocket messages
"""
try:
while self.connection:
message = await self.connection.recv()
print("<<< {}".format(message))
await self.process(json.loads(message))
except websockets.exceptions.ConnectionClosed:
output_msg(out=self.output, status="Websocket Connection closed")
# Commands:::
async def _send(self, no_recv: bool=False, **msg):
"""Send a message out the websocket connection
:param no_recv: Flag to indicate if response is expected
:param msg: message content
:return:
"""
output_msg(out=self.output, flow="output", msg=msg)
try:
await self.connection.send(json.dumps(msg))
if no_recv:
return
message = await self.connection.recv()
await self.process(json.loads(message))
except websockets.exceptions.ConnectionClosed:
pass
async def cmd_connect(self, server: str=None, **kwargs):
"""Connect to a remote websocket server
:param server: Websocket url
:param kwargs: ignored
"""
srv = self.config.server or server
output_msg(out=self.output, status="Connecting to {}".format(srv))
self.connection = await websockets.connect(srv)
self.recv.append(asyncio.ensure_future(self.receiver()))
async def cmd_close(self, **kwargs):
"""Close the websocket connection (if needed)
:param kwargs: ignored
"""
output_msg(out=self.output, status="Closing socket connection")
if self.connection and self.connection.state == 1:
try:
for recv in self.recv:
recv.cancel()
self.recv = []
await self.connection.close()
except (websockets.exceptions.ConnectionClosed,
futures.CancelledError):
pass
async def cmd_sleep(self, period: int=5, **kwargs):
output_msg(out=self.output, status="Sleeping...")
await asyncio.sleep(period)
async def cmd_hello(self, uaid: str=None, **kwargs):
"""Send a websocket "hello" message
:param uaid: User Agent ID (if reconnecting)
"""
if not self.connection or self.connection.state != 1:
await self.cmd_connect()
output_msg(out=self.output, status="Sending Hello")
args = dict(messageType="hello", use_webpush=1, **kwargs)
if uaid:
args['uaid'] = uaid
elif self.uaid:
args['uaid'] = self.uaid
await self._send(**args)
async def cmd_ack(self,
channelID: str=None,
version: str=None,
timeout: int=60,
**kwargs):
"""Acknowledge a previous mesage
:param channelID: Channel to acknowledge
:param version: Version string for message to acknowledge
:param kwargs: Additional optional arguments
:param timeout: Time to wait for notifications (used by testing)
"""
timeout = timeout * 2
while not self.notifications:
output_msg(
out=self.output,
status="No notifications recv'd, Sleeping...")
await asyncio.sleep(0.5)
timeout -= 1
if timeout < 1:
raise PushException("Timeout waiting for messages")
self.notifications.reverse()
for notif in self.notifications:
output_msg(
out=self.output,
status="Sending ACK",
channelID=channelID or notif['channelID'],
version=version or notif['version'])
await self._send(messageType="ack",
channelID=channelID or notif['channelID'],
version=version or notif['version'],
no_recv=True)
self.notifications = []
async def cmd_register(self, channelID: str=None,
key: str=None, **kwargs):
"""Register a new ChannelID
:param channelID: UUID for the channel to register
:param key: applicationServerKey for a restricted access channel
:param kwargs: additional optional arguments
:return:
"""
output_msg(
out=self.output,
status="Sending new channel registration")
channelID = channelID or self.channelID or str(uuid.uuid4())
args = dict(messageType='register',
channelID=channelID)
if key:
args[key] = key
args.update(kwargs)
await self._send(**args)
async def cmd_done(self, **kwargs):
"""Close all connections and mark as done
:param kwargs: ignored
:return:
"""
output_msg(
out=self.output,
status="done")
await self.cmd_close()
await self.connection.close_connection()
"""
recv_* commands handle incoming responses.Since they are asynchronous
and need to trigger follow-up tasks, they each will need to pull and
process the next task.
"""
async def recv_hello(self, **msg: Dict[str, Any]):
"""Process a received "hello"
:param msg: body of response
:return:
"""
assert(msg['status'] == 200)
try:
self.uaid = msg['uaid']
await self._next_task()
except KeyError as ex:
raise PushException from ex
async def recv_register(self, **msg):
"""Process a received registration message
:param msg: body of response
:return:
"""
assert(msg['status'] == 200)
self.pushEndpoint = self._fix_endpoint(msg['pushEndpoint'])
self.channelID = msg['channelID']
output_msg(
out=self.output,
flow="input",
msg=dict(
message="register",
channelID=self.channelID,
pushEndpoint=self.pushEndpoint))
await self._next_task()
async def recv_notification(self, **msg):
"""Process a received notification message.
This event does NOT trigger the next command in the stack.
:param msg: body of response
"""
def repad(string):
return string + '===='[len(msg['data']) % 4:]
msg['_decoded_data'] = base64.urlsafe_b64decode(
repad(msg['data'])).decode()
output_msg(
out=self.output,
flow="input",
message="notification",
msg=msg)
self.notifications.append(msg)
await self.cmd_ack()
async def _post(self, session, url: str, data: bytes):
"""Post a message to the endpoint
:param session: async session object
:param url: pushEndpoint
:param data: data to send
:return:
"""
# print ("Fetching {}".format(url))
with aiohttp.Timeout(10, loop=session.loop):
return await session.post(url=url,
data=data)
async def _post_session(self, url: str,
headers: Dict[str, str],
data: bytes):
"""create a session to send the post message to the endpoint
:param url: pushEndpoint
:param headers: dictionary of headers
:param data: body of the content to send
"""
async with aiohttp.ClientSession(
loop=self.loop,
headers=headers,
read_timeout=30,
connector=self.tls_conn,
) as session:
reply = await self._post(session, url, data)
return reply
async def cmd_push(self, data: bytes=None,
headers: Dict[str, str]=None,
claims: Dict[str, str]=None):
"""Push data to the pushEndpoint
:param data: message content
:param headers: dictionary of headers
:param claims: VAPID claims
:return:
"""
if not self.pushEndpoint:
raise PushException("No Endpoint, no registration?")
if not headers:
headers = {}
if claims:
headers.update(self._cache_sign(claims))
output_msg(
out=self.output,
status="Pushing message",
msg=repr(data))
if data and 'content-encoding' not in headers:
headers.update({
"content-encoding": "aesgcm128",
"encryption": "salt=test",
"encryption-key": "dh=test",
})
result = await self._post_session(self.pushEndpoint, headers, data)
body = await result.text()
output_msg(
out=self.output,
flow="http-out",
pushEndpoint=self.pushEndpoint,
headers=headers,
data=repr(data),
result="{}: {}".format(result.status, body))