| # |
| # |
| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| # |
| |
| import argparse |
| import base64 |
| import functools |
| import json |
| import sys |
| import uuid |
| import weakref |
| |
| from threading import Thread, Lock |
| |
| try: |
| from urllib.parse import urlencode |
| except ImportError: |
| from urllib import urlencode |
| |
| import pika |
| import tornado.websocket |
| import tornado.ioloop |
| import tornado.auth |
| import tornado.escape |
| import tornado.concurrent |
| |
| |
| SETTINGS = {} |
| |
| |
| class Error(Exception): |
| """Base error class for exceptions in this module""" |
| pass |
| |
| class ConsumerConfigError(Error): |
| """Raised when an issue with consumer configuration occurs""" |
| def __init__(self, message): |
| self.message = message |
| |
| class ConsumerKeyError(Error): |
| def __init__(self, message, key): |
| self.message = message |
| self.key = key |
| |
| class AuthError(Error): |
| """Raised when something went wrong during authentication""" |
| def __init__(self, error, code): |
| self.message = error |
| self.code = code |
| |
| |
| |
| class PikaAsyncConsumer(Thread): |
| |
| """ |
| The primary entry point for routing incoming messages to the proper handler. |
| |
| """ |
| |
| def __init__(self, rabbitmq_url, exchange_name, queue_name, |
| exchange_type="direct", routing_key="#"): |
| """ |
| Create a new instance of Streamer. |
| |
| Arguments: |
| rabbitmq_url -- URL to RabbitMQ server |
| exchange_name -- name of RabbitMQ exchange to join |
| queue_name -- name of RabbitMQ queue to join |
| |
| Keyword Arguments: |
| exchange_type -- one of 'direct', 'topic', 'fanout', 'headers' |
| (default 'direct') |
| routing_keys -- the routing key that this consumer listens for |
| (default '#', receives all messages) |
| |
| """ |
| print("Creating new consumer") |
| super(PikaAsyncConsumer, self).__init__(daemon=True) |
| self._connection = None |
| self._channel = None |
| self._shut_down = False |
| self._consumer_tag = None |
| self._url = rabbitmq_url |
| self._client_list = [] |
| self._lock = Lock() |
| |
| # The following are necessary to guarantee that both the RabbitMQ |
| # server and Streamer know where to look for messages. These names will |
| # be decided before dispatch and should be recorded in a config file or |
| # else on a per-job basis. |
| self._exchange = exchange_name |
| self._exchange_type = exchange_type |
| self._queue = queue_name |
| self._routing_key = routing_key |
| |
| def add_client(self, client): |
| """Add a new client to the recipient list. |
| |
| Arguments: |
| client -- a reference to the client object to add |
| """ |
| self._lock.acquire() |
| # Create a weakref to ensure that cyclic references to WebSocketHandler |
| # objects do not cause problems for garbage collection |
| self._client_list.append(weakref.ref(client)) |
| self._lock.release() |
| |
| def remove_client(self, client): |
| """Remove a client from the recipient list. |
| |
| Arguments: |
| client -- a reference to the client object to remove |
| """ |
| self._lock.acquire() |
| for i in range(0, len(self._client_list)): |
| # Parentheses after _client_list[i] to deference the weakref to its |
| # strong reference |
| if self._client_list[i]() is client: |
| self._client_list.pop(i) |
| break |
| self._lock.release() |
| |
| |
| def connect(self): |
| """ |
| Create an asynchronous connection to the RabbitMQ server at URL. |
| |
| """ |
| return pika.SelectConnection(pika.URLParameters(self._url), |
| on_open_callback=self.on_connection_open, |
| on_close_callback=self.on_connection_close, |
| stop_ioloop_on_close=False) |
| |
| def on_connection_open(self, unused_connection): |
| """ |
| Actions to perform when the connection opens. This may not happen |
| immediately, so defer action to this callback. |
| |
| Arguments: |
| unused_connection -- the created connection (by this point already |
| available as self._connection) |
| |
| """ |
| self._connection.channel(on_open_callback=self.on_channel_open) |
| |
| def on_connection_close(self, connection, code, text): |
| """ |
| Actions to perform when the connection is unexpectedly closed by the |
| RabbitMQ server. |
| |
| Arguments: |
| connection -- the connection that was closed (same as self._connection) |
| code -- response code from the RabbitMQ server |
| text -- response body from the RabbitMQ server |
| |
| """ |
| self._channel = None |
| if self._shut_down: |
| self._connection.ioloop.stop() |
| else: |
| self._connection.add_timeout(5, self.reconnect) |
| |
| def reconnect(self): |
| """ |
| Attempt to reestablish a connection with the RabbitMQ server. |
| """ |
| self._connection.ioloop.stop() # Stop the ioloop to completely close |
| |
| if not self._shut_down: # Connect and restart the ioloop |
| self._connection = self.connect() |
| self._connection.ioloop.start() |
| |
| def on_channel_open(self, channel): |
| """ |
| Store the opened channel for future use and set up the exchange and |
| queue to be used. |
| |
| Arguments: |
| channel -- the Channel instance opened by the Channel.Open RPC |
| """ |
| self._channel = channel |
| self._channel.add_on_close_callback(self.on_channel_close) |
| self.declare_exchange() |
| |
| |
| def on_channel_close(self, channel, code, text): |
| """ |
| Actions to perform when the channel is unexpectedly closed by the |
| RabbitMQ server. |
| |
| Arguments: |
| connection -- the connection that was closed (same as self._connection) |
| code -- response code from the RabbitMQ server |
| text -- response body from the RabbitMQ server |
| """ |
| self._connection.close() |
| |
| def declare_exchange(self): |
| """ |
| Set up the exchange that will route messages to this consumer. Each |
| RabbitMQ exchange is uniquely identified by its name, so it does not |
| matter if the exchange has already been declared. |
| """ |
| self._channel.exchange_declare(self.declare_exchange_success, |
| self._exchange, |
| self._exchange_type) |
| |
| def declare_exchange_success(self, unused_connection): |
| """ |
| Actions to perform on successful exchange declaration. |
| """ |
| self.declare_queue() |
| |
| def declare_queue(self): |
| """ |
| Set up the queue that will route messages to this consumer. Each |
| RabbitMQ queue can be defined with routing keys to use only one |
| queue for multiple jobs. |
| """ |
| self._channel.queue_declare(self.declare_queue_success, |
| self._queue) |
| |
| def declare_queue_success(self, method_frame): |
| """ |
| Actions to perform on successful queue declaration. |
| """ |
| self._channel.queue_bind(self.munch, |
| self._queue, |
| self._exchange, |
| self._routing_key |
| ) |
| |
| def munch(self, unused): |
| """ |
| Begin consuming messages from the Airavata API server. |
| """ |
| self._channel.add_on_cancel_callback(self.cancel_channel) |
| self._consumer_tag = self._channel.basic_consume(self._process_message) |
| |
| def cancel_channel(self, method_frame): |
| if self._channel is not None: |
| self._channel._close() |
| |
| def _process_message(self, ch, method, properties, body): |
| """ |
| Receive and verify a message, then pass it to the router. |
| |
| Arguments: |
| ch -- the channel that routed the message |
| method -- delivery information |
| properties -- message properties |
| body -- the message |
| """ |
| print("Received Message: %s" % body) |
| self._lock.acquire() |
| for client in self._client_list: |
| # Parentheses after client to deference the weakref to its |
| # strong reference |
| client().write_message(body) |
| self._lock.release() |
| self._channel.basic_ack(delivery_tag=method.delivery_tag) |
| |
| def stop_consuming(self): |
| """ |
| Stop the consumer if active. |
| """ |
| if self._channel: |
| self._channel.basic_cancel(self.close_channel, self._consumer_tag) |
| |
| def close_channel(self, unused): |
| """ |
| Close the channel to shut down the consumer and connection. |
| """ |
| self._channel.queue_delete(queue=self._queue) |
| self._channel.close() |
| |
| def run(self): |
| """ |
| Start a connection with the RabbitMQ server. |
| """ |
| self._connection = self.connect() |
| self._connection.ioloop.start() |
| |
| def stop(self): |
| """ |
| Stop an active connection with the RabbitMQ server. |
| """ |
| self._closing = True |
| self.stop_consuming() |
| |
| |
| class Wso2OAuth2Mixin(tornado.auth.OAuth2Mixin): |
| _OAUTH_AUTHORIZE_URL = "https://idp.scigap.org:9443/oauth2/authorize" |
| _OAUTH_ACCESS_TOKEN_URL = "https://idp.scigap.org:9443/oauth2/token" |
| |
| @tornado.auth._auth_return_future |
| def get_authenticated_user(self, username, password, callback=None): |
| print("Authenticating user %s" % (username)) |
| http = self.get_auth_http_client() |
| body = urlencode({ |
| "client_id": SETTINGS["oauth_client_key"], |
| "client_secret": SETTINGS["oauth_client_secret"], |
| "grant_type": SETTINGS["oauth_grant_type"], |
| "username": username, |
| "password": password |
| }) |
| http.fetch(self._OAUTH_ACCESS_TOKEN_URL, functools.partial(self._on_access_token, callback), method="POST", body=body) |
| |
| def _on_access_token(self, future, response): |
| if response.error: |
| print(str(response)) |
| print(response.body) |
| print(response.error) |
| future.set_exception(AuthError(response.error, response.code)) |
| return |
| |
| print(response.body) |
| future.set_result(tornado.escape.json_decode(response.body)) |
| |
| class AuthHandler(tornado.web.RequestHandler, Wso2OAuth2Mixin): |
| def get_current_user(self): |
| expires_in = self.get_secure_cookie("expires-in", max_age_days=SETTINGS['maximum_cookie_age']) |
| print(expires_in) |
| if expires_in: |
| return self.get_secure_cookie("ws-auth-token", max_age_days=float(expires_in)) |
| return None |
| |
| def set_default_headers(self): |
| self.set_header("Content-Type", "text/plain") |
| self.set_header("Access-Control-Allow-Origin", "*") |
| self.set_header("Access-Control-Allow-Headers", "x-requested-with") |
| self.set_header('Access-Control-Allow-Methods', 'POST, GET, OPTIONS') |
| |
| def get(self): |
| if self.get_current_user(): |
| self.set_status(200) |
| print("Authenticated") |
| self.write("Authenticated") |
| |
| else: |
| self.set_status(403) |
| print("Not Authenticated") |
| self.write("Not Authenticated") |
| |
| @tornado.gen.coroutine |
| def post(self): |
| try: |
| username = self.get_body_argument("username") |
| password = self.get_body_argument("password") |
| redirect = self.get_body_argument("redirect") |
| if username == "" or password == "": |
| raise tornado.web.MissingArgumentError |
| |
| access = yield self.get_authenticated_user(username, password) |
| days = (access["expires_in"] / 3600) / 24 # Convert to days |
| print(days) |
| self.set_secure_cookie("ws-auth-token", |
| access["access_token"], |
| expires_days=days) |
| self.set_secure_cookie("expires-in", |
| str(1), |
| expires_days=SETTINGS['maximum_cookie_age']) |
| self.write("Success") |
| except tornado.web.MissingArgumentError: |
| print("Missing an argument") |
| self.set_status(400) |
| self.write("Authentication information missing") |
| except AuthError as e: |
| print("The future freaks me out") |
| self.set_status(access.code) |
| self.set_header("Content-Type", "text/html") |
| self.write(access.message) |
| |
| success_code = """<p>Redirecting to <a href="%(url)s">%(url)s</a></p> |
| <script type="text/javascript"> |
| window.location = %(url)s; |
| </script> |
| """ % { 'url': redirect} |
| self.set_status(200) |
| self.redirect(redirect) |
| #return self.render_string(success_code) |
| |
| |
| |
| class AMQPWSHandler(tornado.websocket.WebSocketHandler):#, Wso2OAuth2Mixin): |
| |
| """ |
| Pass messages to a connected WebSockets client. |
| |
| A subclass of the Tornado WebSocketHandler class, this class takes no |
| action when receiving a message from the client. Instead, it is associated |
| with an AMQP consumer and writes a message to the client each time one is |
| consumed in the queue. |
| """ |
| |
| # def set_default_headers(self): |
| # self.set_header("Access-Control-Allow-Origin", "*") |
| # self.set_header("Access-Control-Allow-Headers", "x-requested-with") |
| # self.set_header('Access-Control-Allow-Methods', 'POST, GET, OPTIONS') |
| |
| def check_origin(self, origin): |
| """Check the domain origin of the connection request. |
| |
| This can be made more robust to ensure that connections are only |
| accepted from verified PGAs. |
| |
| Arguments: |
| origin -- the value of the Origin HTTP header |
| """ |
| return True |
| |
| def open(self, resource_type, resource_id): |
| """Associate a new connection with a consumer. |
| |
| When a new connection is opened, it is a request to retrieve data |
| from an AMQP queue. The open operation should also do some kind of |
| authentication. |
| |
| Arguments: |
| resource_type -- "experiment" or "project" or "data" |
| resource_id -- the Airavata id for the resource |
| """ |
| self.stream.set_nodelay(True) |
| self.resource_id = resource_id |
| self.write_message("Opened the connection") |
| |
| self.add_to_consumer() |
| |
| # expires_in = self.get_secure_cookie("expires_in", max_age_days=SETTINGS["maximum_cookie_age"]) |
| # if expires_in is not None and self.get_secure_cookie("ws-auth-token", max_age_days=float(expires_in)): |
| # print("Found secure cookie") |
| # self.write_message("Authenticated") |
| # self.add_to_consumer() |
| # else: |
| # print("Closing connection") |
| # self.close() |
| |
| def on_message(self, message): |
| """Handle incoming messages from the client. |
| |
| Tornado requires subclasses to override this method, however in this |
| case we do not wish to take any action when receiving a message from |
| the client. The purpose of this class is only to push messages to the |
| client. |
| """ |
| print(message) |
| message = tornado.escape.json_decode(message) |
| access = yield self.get_authenticated_user(message["username"], message["password"]) |
| access = access |
| days = (access["expires_in"] / 3600) / 24 # Convert to days |
| print(days) |
| self.set_secure_cookie("ws-auth-token", |
| access["access_token"], |
| expires_days=days) |
| self.set_secure_cookie("expires_in", |
| str(days), |
| expires_days=SETTINGS['maximum_cookie_age']) |
| |
| |
| def on_close(self): |
| try: |
| print("Closing connection") |
| self.application.remove_client_from_consumer(self.resource_id, self) |
| except KeyError: |
| print("Error: resource %s does not exist" % self.resource_id) |
| finally: |
| self.close() |
| |
| def add_to_consumer(self): |
| try: |
| self.application.add_client_to_consumer(self.resource_id, self) |
| except AttributeError as e: |
| print("Error: tornado.web.Application object is not AMQPWSTunnel") |
| print(e) |
| |
| |
| class AMQPWSTunnel(tornado.web.Application): |
| |
| """ |
| Send messages from an AMQP queue to WebSockets clients. |
| |
| In addition to the standard Tornado Application class functionality, this |
| class maintains a list of active AMQP consumers and maps WebSocketHandlers |
| to the correct consumers. |
| """ |
| |
| def __init__(self, consumer_list=None, consumer_config=None, handlers=None, |
| default_host='', transforms=None, **settings): |
| print("Starting AMQP-WS-Tunnel application") |
| super(AMQPWSTunnel, self).__init__(handlers=handlers, |
| default_host=default_host, |
| transforms=transforms, |
| **settings) |
| |
| self.consumer_list = {} if consumer_list is None else consumer_list |
| if consumer_config is None: |
| raise ConsumerConfigError("No consumer configuration provided") |
| self.consumer_config = consumer_config |
| |
| def consumer_exists(self, resource_id): |
| """Determine if a consumer exists for a particular resource. |
| |
| Arguments: |
| resource_id -- the consumer to find |
| """ |
| return resource_id in self.consumer_list |
| |
| def add_client_to_consumer(self, resource_id, client): |
| """Add a new client to a consumer's messaging list. |
| |
| Arguments: |
| resource_id -- the consumer to add to |
| client -- the client to add |
| """ |
| if not self.consumer_exists(resource_id): |
| print("Creating new consumer") |
| print(self.consumer_config) |
| consumer = PikaAsyncConsumer(self.consumer_config["rabbitmq_url"], |
| self.consumer_config["exchange_name"], |
| self.consumer_config["queue_name"], |
| exchange_type=self.consumer_config["exchange_type"], |
| routing_key=resource_id) |
| print("Adding to consumer list") |
| self.consumer_list[resource_id] = consumer |
| print("Starting consumer") |
| consumer.start() |
| |
| print("Adding new client to %s" % (resource_id)) |
| consumer = self.consumer_list[resource_id] |
| consumer.add_client(client) |
| |
| def remove_client_from_consumer(self, resource_id, client): |
| """Remove a client from a consumer's messaging list. |
| |
| Arguments: |
| resource_id -- the consumer to remove from |
| client -- the client to remove |
| """ |
| if self.consumer_exists(resource_id): |
| print("Removing client from %s" % (resource_id)) |
| self.consumer_list[resource_id].remove_client(client) |
| #else: |
| # raise ConsumerKeyError("Trying to remove client from nonexistent consumer", resource_id) |
| |
| def shutdown(self): |
| """Shut down the application and release all resources. |
| |
| |
| """ |
| for name, consumer in self.consumer_list.items(): |
| consumer.stop() |
| #consumer.join() |
| #self.consumer_list[name] = None |
| |
| #self.consumer_list = {} |
| |
| |
| |
| if __name__ == "__main__": |
| i = open(sys.argv[1]) |
| config = json.load(i) |
| i.close() |
| |
| SETTINGS["oauth_client_key"] = config["oauth_client_key"] |
| SETTINGS["oauth_client_secret"] = config["oauth_client_secret"] |
| SETTINGS["oauth_grant_type"] = config["oauth_grant_type"] |
| SETTINGS["maximum_cookie_age"] = config["maximum_cookie_age"] |
| |
| settings = { |
| "cookie_secret": base64.b64encode(uuid.uuid4().bytes + uuid.uuid4().bytes), |
| #"xsrf_cookies": True |
| } |
| |
| application = AMQPWSTunnel(handlers=[ |
| (r"/auth", AuthHandler), |
| (r"/(experiment)/(.+)", AMQPWSHandler) |
| ], |
| consumer_config=config, |
| debug=True, |
| **settings) |
| |
| application.listen(8888) |
| |
| try: |
| tornado.ioloop.IOLoop.current().start() |
| except KeyboardInterrupt: |
| application.shutdown() |
