localwsbroadcastserver.py

# -*- coding: utf-8 -*-
"""
    :description: A Function To Generate TOTP From Qrcode Image.
    :license: MIT.
    :author: Shabbir Hasan
    :created: On Wednesday November 18 2022 17:43:57 GMT+05:30
"""
__author__ = "Shabbir Hasan"
__webpage__ = "https://github.com/ShabbirHasan1"
__license__ = "MIT"


import json
import pytz
from datetime import datetime
from twisted.internet.protocol import ReconnectingClientFactory
from autobahn.twisted.websocket import (
    WebSocketServerFactory,
    WebSocketServerProtocol,
    WebSocketClientFactory,
    WebSocketClientProtocol,
)


class DateTimeEncoder(json.JSONEncoder):
    def default(self, o):
        if isinstance(o, datetime):
            return {"TimeStr": o.strftime("%Y-%m-%dT%H:%M:%S.%f%z")}
        return json.JSONEncoder.default(self, o)


class DateTimeDecoder(json.JSONDecoder):
    def default(self, o):
        TimeStr = o.get("TimeStr")
        if TimeStr is not None:
            try:
                return datetime.strptime(TimeStr, "%Y-%m-%dT%H:%M:%S.%f%z")
            except ValueError:
                return datetime.strptime(TimeStr, "%Y-%m-%dT%H:%M:%S.%f")
        return json.JSONDecoder.default(self, o)


class BroadcastClientProtocol(WebSocketClientProtocol):
    def onMessage(self, payload, isBinary):
        if not isBinary:
            print(
                "Tick Received: {}".format(
                    json.loads(payload.decode("utf8"), cls=DateTimeDecoder)
                )
            )


class BroadcastServerProtocol(WebSocketServerProtocol):
    def onOpen(self):
        self.factory.register(self)

    def onMessage(self, payload, isBinary):
        if not isBinary:
            msg = "{} from {}".format(payload.decode("utf8"), self.peer)
            self.factory.broadcast(msg)

    def connectionLost(self, reason):
        WebSocketServerProtocol.connectionLost(self, reason)
        self.factory.unregister(self)


class BroadcastServerFactory(WebSocketServerFactory, ReconnectingClientFactory):
    def __init__(self, ws_url):
        self.clients = []
        WebSocketServerFactory.__init__(self, ws_url)

    def register(self, client):
        if client not in self.clients:
            print("registered client {}".format(client.peer))
            self.clients.append(client)

    def unregister(self, client):
        if client in self.clients:
            print("unregistered client {}".format(client.peer))
            self.clients.remove(client)

    def broadcast(self, msg, optimized=True):
        if optimized:
            print("broadcasting '{}' ..".format(msg))
            preparedMsg = self.prepareMessage(msg.encode("utf8"))
            [c.sendPreparedMessage(preparedMsg) for c in self.clients]
        else:
            print("broadcasting '{}' ..".format(msg))
            [c.sendMessage(msg.encode("utf8")) for c in self. Clients]