pypubsub/pypubsub.py

#!/usr/bin/env python3
# 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.

"""PyPubSub - a simple publisher/subscriber service written in Python 3"""
import asyncio
import aiohttp.web
import aiofile
import os
import time
import json
import yaml
import netaddr
import binascii
import base64
import argparse
import collections
import plugins.ldap
import plugins.sqs
import typing

# Some consts
PUBSUB_VERSION = '0.7.1'
PUBSUB_CONTENT_TYPE = 'application/vnd.pypubsub-stream'
PUBSUB_DEFAULT_PORT = 2069
PUBSUB_DEFAULT_IP = '0.0.0.0'
PUBSUB_DEFAULT_MAX_PAYLOAD_SIZE = 102400
PUBSUB_DEFAULT_BACKLOG_SIZE = 0
PUBSUB_DEFAULT_BACKLOG_AGE = 0
PUBSUB_BAD_REQUEST = "I could not understand your request, sorry! Please see https://pubsub.apache.org/api.html \
for usage documentation.\n"
PUBSUB_PAYLOAD_RECEIVED = "Payload received, thank you very much!\n"
PUBSUB_NOT_ALLOWED = "You are not authorized to deliver payloads!\n"
PUBSUB_BAD_PAYLOAD = "Bad payload type. Payloads must be JSON dictionary objects, {..}!\n"
PUBSUB_PAYLOAD_TOO_LARGE = "Payload is too large for me to serve, please make it shorter.\n"
PUBSUB_WRITE_TIMEOUT = 0.35  # If we can't deliver to a pipe within N seconds, drop it.


class ServerConfig(typing.NamedTuple):
    ip: str
    port: int
    payload_limit: int


class BacklogConfig(typing.NamedTuple):
    max_age: int
    queue_size: int
    storage: typing.Optional[str]


class Configuration:
    server: ServerConfig
    backlog: BacklogConfig
    payloaders: typing.List[netaddr.ip.IPNetwork]
    oldschoolers: typing.List[str]

    def __init__(self, yml: dict):

        # LDAP Settings
        self.ldap = None
        lyml = yml.get('clients', {}).get('ldap')
        if isinstance(lyml, dict):
            self.ldap = plugins.ldap.LDAPConnection(lyml)

        # SQS?
        self.sqs = yml.get('sqs')

        # Main server config
        server_ip = yml['server'].get('bind', PUBSUB_DEFAULT_IP)
        server_port = int(yml['server'].get('port', PUBSUB_DEFAULT_PORT))
        server_payload_limit = int(yml['server'].get('max_payload_size', PUBSUB_DEFAULT_MAX_PAYLOAD_SIZE))
        self.server = ServerConfig(ip=server_ip, port=server_port, payload_limit=server_payload_limit)

        # Backlog settings
        bma = yml['server'].get('backlog', {}).get('max_age', PUBSUB_DEFAULT_BACKLOG_AGE)
        if isinstance(bma, str):
            bma = bma.lower()
            if bma.endswith('s'):
                bma = int(bma.replace('s', ''))
            elif bma.endswith('m'):
                bma = int(bma.replace('m', '')) * 60
            elif bma.endswith('h'):
                bma = int(bma.replace('h', '')) * 3600
            elif bma.endswith('d'):
                bma = int(bma.replace('d', '')) * 86400
        bqs = yml['server'].get('backlog', {}).get('size',
                                                   PUBSUB_DEFAULT_BACKLOG_SIZE)
        bst = yml['server'].get('backlog', {}).get('storage')
        self.backlog = BacklogConfig(max_age=bma, queue_size=bqs, storage=bst)

        # Payloaders - clients that can post payloads
        self.payloaders = [netaddr.IPNetwork(x) for x in yml['clients'].get('payloaders', [])]

        # Binary backwards compatibility
        self.oldschoolers = yml['clients'].get('oldschoolers', [])


class Server:
    """Main server class, responsible for handling requests and publishing events """
    yaml: dict
    config: Configuration
    subscribers: list
    pending_events: asyncio.Queue
    backlog: list
    last_ping = typing.Type[float]
    server: aiohttp.web.Server

    def __init__(self, args: argparse.Namespace):
        self.yaml = yaml.safe_load(open(args.config))
        self.config = Configuration(self.yaml)
        self.subscribers = []
        self.pending_events = asyncio.Queue()
        self.backlog = []
        self.last_ping = time.time()
        self.acl = {}
        try:
            self.acl = yaml.safe_load(open(args.acl))
        except FileNotFoundError:
            print(f"ACL configuration file {args.acl} not found, private events will not be broadcast.")

    async def poll(self):
        """Polls for new stuff to publish, and if found, publishes to whomever wants it."""
        while True:
            payload: Payload = await self.pending_events.get()
            bad_subs: list = await payload.publish(self.subscribers)
            self.pending_events.task_done()

            # Cull subscribers we couldn't deliver payload to.
            for bad_sub in bad_subs:
                print("Culling %r due to connection errors" % bad_sub)
                try:
                    self.subscribers.remove(bad_sub)
                except ValueError:  # Already removed elsewhere
                    pass

    async def handle_request(self, request: aiohttp.web.BaseRequest):
        """Generic handler for all incoming HTTP requests"""
        resp: typing.Union[aiohttp.web.Response, aiohttp.web.StreamResponse]

        # Define response headers first...
        headers = {
            'Server': 'PyPubSub/%s' % PUBSUB_VERSION,
            'X-Subscribers': str(len(self.subscribers)),
            'X-Requests': str(self.server.requests_count),
        }

        # Are we handling a publisher payload request? (PUT/POST)
        if request.method in ['PUT', 'POST']:
            ip = netaddr.IPAddress(request.remote)
            allowed = False
            for network in self.config.payloaders:
                if ip in network:
                    allowed = True
                    break
            if not allowed:
                resp = aiohttp.web.Response(headers=headers, status=403, text=PUBSUB_NOT_ALLOWED)
                return resp
            if request.can_read_body:
                try:
                    if request.content_length and request.content_length > self.config.server.payload_limit:
                        resp = aiohttp.web.Response(headers=headers, status=400, text=PUBSUB_PAYLOAD_TOO_LARGE)
                        return resp
                    body = await request.text()
                    as_json = json.loads(body)
                    assert isinstance(as_json, dict)  # Payload MUST be an dictionary object, {...}
                    pl = Payload(request.path, as_json)
                    self.pending_events.put_nowait(pl)
                    # Add to backlog?
                    if self.config.backlog.queue_size > 0:
                        self.backlog.append(pl)
                        # If backlog has grown too large, delete the first (oldest) item in it.
                        while len(self.backlog) > self.config.backlog.queue_size:
                            del self.backlog[0]

                    resp = aiohttp.web.Response(headers=headers, status=202, text=PUBSUB_PAYLOAD_RECEIVED)
                    return resp
                except json.decoder.JSONDecodeError:
                    resp = aiohttp.web.Response(headers=headers, status=400, text=PUBSUB_BAD_REQUEST)
                    return resp
                except AssertionError:
                    resp = aiohttp.web.Response(headers=headers, status=400, text=PUBSUB_BAD_PAYLOAD)
                    return resp
        # Is this a subscriber request? (GET)
        elif request.method == 'GET':
            resp = aiohttp.web.StreamResponse(headers=headers)
            # We do not support HTTP 1.0 here...
            if request.version.major == 1 and request.version.minor == 0:
                return resp
            subscriber = Subscriber(self, resp, request)

            # Is there a basic auth in this request? If so, set up ACL
            auth = request.headers.get('Authorization')
            if auth:
                subscriber.acl = await subscriber.parse_acl(auth)

            # Subscribe the user before we deal with the potential backlog request and pings
            self.subscribers.append(subscriber)
            resp.content_type = PUBSUB_CONTENT_TYPE
            try:
                resp.enable_chunked_encoding()
                await resp.prepare(request)

                # Is the client requesting a backlog of items?
                backlog = request.headers.get('X-Fetch-Since')
                if backlog:
                    try:
                        backlog_ts = int(backlog)
                    except ValueError:  # Default to 0 if we can't parse the epoch
                        backlog_ts = 0
                    # If max age is specified, force the TS to minimum that age
                    if self.config.backlog.max_age > 0:
                        backlog_ts = max(backlog_ts, int(time.time() - self.config.backlog.max_age))
                    # For each item, publish to client if new enough.
                    for item in self.backlog:
                        if item.timestamp >= backlog_ts:
                            await item.publish([subscriber])

                while True:
                    await subscriber.ping()
                    if subscriber not in self.subscribers:  # If we got dislodged somehow, end session
                        break
                    await asyncio.sleep(5)
            # We may get exception types we don't have imported, so grab ANY exception and kick out the subscriber
            except:
                pass
            if subscriber in self.subscribers:
                self.subscribers.remove(subscriber)
            return resp
        elif request.method == 'HEAD':
            resp = aiohttp.web.Response(headers=headers, status=204, text="")
            return resp
        #  I don't know this type of request :/ (DELETE, PATCH, etc)
        else:
            resp = aiohttp.web.Response(headers=headers, status=400, text=PUBSUB_BAD_REQUEST)
            return resp

    async def write_backlog_storage(self):
        previous_backlog = []
        while True:
            if self.config.backlog.storage:
                try:
                    backlog_list = self.backlog.copy()
                    if backlog_list != previous_backlog:
                        previous_backlog = backlog_list
                        async with aiofile.AIOFile(self.config.backlog.storage, 'w+') as afp:
                            offset = 0
                            for item in backlog_list:
                                js =json.dumps({
                                    'timestamp': item.timestamp,
                                    'topics': item.topics,
                                    'json': item.json,
                                    'private': item.private
                                }) + '\n'
                                await afp.write(js, offset=offset)
                                offset += len(js)
                            await afp.fsync()
                except Exception as e:
                    print(f"Could not write to backlog file {self.config.backlog.storage}: {e}")
            await asyncio.sleep(10)

    def read_backlog_storage(self):
        if self.config.backlog.storage and os.path.exists(self.config.backlog.storage):
            try:
                readlines = 0
                with open(self.config.backlog.storage, 'r') as fp:
                    for line in fp.readlines():
                        js = json.loads(line)
                        readlines += 1
                        ppath = "/".join(js['topics'])
                        if js['private']:
                            ppath = '/private/' + ppath
                        payload = Payload(ppath, js['json'], js['timestamp'])
                        self.backlog.append(payload)
                        if self.config.backlog.queue_size < len(self.backlog):
                            self.backlog.pop(0)
            except Exception as e:
                print(f"Error while reading backlog: {e}")

            print(f"Read {readlines} objects from {self.config.backlog.storage}, applied {len(self.backlog)} to backlog.")

    async def server_loop(self, loop: asyncio.BaseEventLoop):
        self.server = aiohttp.web.Server(self.handle_request)
        runner = aiohttp.web.ServerRunner(self.server)
        await runner.setup()
        site = aiohttp.web.TCPSite(runner, self.config.server.ip, self.config.server.port)
        await site.start()
        print("==== PyPubSub v/%s starting... ====" % PUBSUB_VERSION)
        print("==== Serving up PubSub goodness at %s:%s ====" % (
            self.config.server.ip, self.config.server.port))
        if self.config.sqs:
            for key, config in self.config.sqs.items():
                loop.create_task(plugins.sqs.get_payloads(self, config))
        self.read_backlog_storage()
        loop.create_task(self.write_backlog_storage())
        await self.poll()

    def run(self):
        loop = asyncio.get_event_loop()
        try:
            loop.run_until_complete(self.server_loop(loop))
        except KeyboardInterrupt:
            pass
        loop.close()


class Subscriber:
    """Basic subscriber (client) class. Holds information about the connection and ACL"""
    acl:    dict
    topics: typing.List[typing.List[str]]

    def __init__(self, server: Server, connection: aiohttp.web.StreamResponse, request: aiohttp.web.BaseRequest):
        self.connection = connection
        self.acl = {}
        self.server = server
        self.lock = asyncio.Lock()

        # Set topics subscribed to
        self.topics = []
        for topic_batch in request.path.split(','):
            sub_to = [x for x in topic_batch.split('/') if x]
            self.topics.append(sub_to)

        # Is the client old and expecting zero-terminators?
        self.old_school = False
        for ua in self.server.config.oldschoolers:
            if ua in request.headers.get('User-Agent', ''):
                self.old_school = True
                break

    async def parse_acl(self, basic: str):
        """Sets the ACL if possible, based on Basic Auth"""
        try:
            decoded = str(base64.decodebytes(bytes(basic.replace('Basic ', ''), 'ascii')), 'utf-8')
            u, p = decoded.split(':', 1)
            if u in self.server.acl:
                acl_pass = self.server.acl[u].get('password')
                if acl_pass and acl_pass == p:
                    acl = self.server.acl[u].get('acl', {})
                    # Vet ACL for user
                    assert isinstance(acl, dict), f"ACL for user {u} " \
                                                  f"must be a dictionary of sub-IDs and topics, but is not."
                    # Make sure each ACL segment is a list of topics
                    for k, v in acl.items():
                        assert isinstance(v, list), f"ACL segment {k} for user {u} is not a list of topics!"
                    print(f"Client {u} successfully authenticated (and ACL is valid).")
                    return acl
            elif self.server.config.ldap:
                acl = {}
                groups = await self.server.config.ldap.get_groups(u,p)
                # Make sure each ACL segment is a list of topics
                for k, v in self.server.config.ldap.acl.items():
                    if k in groups:
                        assert isinstance(v, dict), f"ACL segment {k} for user {u} is not a dictionary of segments!"
                        for segment, topics in v.items():
                            print(f"Enabling ACL segment {segment} for user {u}")
                            assert isinstance(topics,
                                              list), f"ACL segment {segment} for user {u} is not a list of topics!"
                            acl[segment] = topics
                return acl
        except binascii.Error as e:
            pass  # Bad Basic Auth params, bail quietly
        except AssertionError as e:
            print(e)
            print(f"ACL configuration error: ACL scheme for {u} contains errors, setting ACL to nothing.")
        except Exception as e:
            print(f"Basic unknown exception occurred: {e}")
        return {}

    async def ping(self):
        """Generic ping-back to the client"""
        js = b"%s\n" % json.dumps({"stillalive": time.time()}).encode('utf-8')
        if self.old_school:
            js += b"\0"
        async with self.lock:
            await asyncio.wait_for(self.connection.write(js), timeout=PUBSUB_WRITE_TIMEOUT)


class Payload:
    """A payload (event) object sent by a registered publisher."""

    def __init__(self, path: str, data: dict, timestamp: typing.Optional[float] = None):
        self.json = data
        self.timestamp = timestamp or time.time()
        self.topics = [x for x in path.split('/') if x]
        self.private = False

        # Private payload?
        if self.topics and self.topics[0] == 'private':
            self.private = True
            del self.topics[0]  # Remove the private bit from topics now.

        self.json['pubsub_timestamp'] = self.timestamp
        self.json['pubsub_topics'] = self.topics
        self.json['pubsub_path'] = path

    async def publish(self, subscribers: typing.List[Subscriber]):
        """Publishes an object to all subscribers using those topics (or a sub-set thereof)"""
        js = b"%s\n" % json.dumps(self.json).encode('utf-8')
        ojs = js + b"\0"
        bad_subs = []
        for sub in subscribers:
            # If a private payload, check ACL and bail if not a match
            if self.private:
                can_see = False
                for key, private_topics in sub.acl.items():
                    if all(el in self.topics for el in private_topics):
                        can_see = True
                        break
                if not can_see:
                    continue
            # If subscribed to all the topics, tell a subscriber about this
            for topic_batch in sub.topics:
                if all(el in self.topics for el in topic_batch):
                    try:
                        if sub.old_school:
                            async with sub.lock:
                                await asyncio.wait_for(sub.connection.write(ojs), timeout=PUBSUB_WRITE_TIMEOUT)
                        else:
                            async with sub.lock:
                                await asyncio.wait_for(sub.connection.write(js), timeout=PUBSUB_WRITE_TIMEOUT)
                    except Exception:
                        bad_subs.append(sub)
                    break
        return bad_subs


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", help="Configuration file to load (default: pypubsub.yaml)", default="pypubsub.yaml")
    parser.add_argument("--acl", help="ACL Configuration file to load (default: pypubsub_acl.yaml)",
                        default="pypubsub_acl.yaml")
    cliargs = parser.parse_args()
    pubsub_server = Server(cliargs)
    pubsub_server.run()