bin.bgpout

  1"""
  2Ingest all records, queue, reformat, send, publish, save them.
  3"""
  4
  5__all__ = ["BGPOut"]
  6
  7import os
  8import sys
  9import json
 10from typing import TextIO
 11from Databases.database import BGPDatabases
 12from bgpgraph import BGPGraph
 13
 14
 15class SetEncoder(json.JSONEncoder):
 16    def default(self, obj):
 17        if isinstance(obj, set):
 18            return list(obj)
 19        return json.JSONEncoder.default(self, obj)
 20
 21
 22class BGPOut:
 23    """
 24    BGP Output
 25    - Ingest records, convert them to json, print to console and/or write them to files.
 26    - Save all records in database.
 27    - Can store bgp records in queue before processing them if desired.
 28    """
 29
 30    def __init__(self) -> None:
 31        self.__expected_result = None
 32        self.__json_out = None
 33        self.verbose: bool = False
 34        """print to console"""
 35        self.isQueue: bool = False
 36        """Enable queue, can prevent from blocking BGPStream"""
 37        self.isStarted: bool = False
 38        """Is the stream started or not"""
 39        self.databases = BGPDatabases({})
 40        self.graph = BGPGraph()
 41
 42    #######################
 43    #   GETTERS/SETTERS   #
 44    #######################
 45
 46    @property
 47    def json_out(self) -> TextIO:
 48        """
 49        (File): Print JSON in it
 50
 51        Raises:
 52            Exception: If unable write/read in it
 53        """
 54        return self.__json_out
 55
 56    @json_out.setter
 57    def json_out(self, json_out):
 58        """
 59        Setter for JSON output
 60        Parameters:
 61            json_out (File): Where to output json
 62
 63        Raises:
 64            Exception: If unable write/read in it
 65        """
 66        if hasattr(json_out, "write"):
 67            self.__json_out = json_out
 68
 69    @property
 70    def expected_result(self) -> TextIO:
 71        """Expected result when execution end"""
 72        return self.__expected_result
 73
 74    @expected_result.setter
 75    def expected_result(self, exp):
 76        if exp is not None:
 77            if hasattr(exp, "read"):
 78                self.__expected_result = exp
 79            else:
 80                raise FileNotFoundError(f"Is {exp} a file ?")
 81
 82    ########
 83    # MAIN #
 84    ########
 85
 86    def start(self):
 87        """
 88        Start output
 89
 90        - Init json output if specified
 91        - Start queue if enabled
 92        """
 93        if self.__json_out:
 94            self.__json_out.write("[")
 95        self.isStarted = True
 96
 97    def stop(self):
 98        """
 99        - Set state as stopped
100        - Close file output
101        - Check if result is as expected
102        """
103        if self.isStarted:
104            self.isStarted = False
105            self.databases.stop()
106            self.closeFile(self.__json_out)
107            if self.__expected_result:
108                print("Testing result: ")
109                with open(self.__json_out.name, "r+") as js:
110                    print(
111                        "Result is as expected"
112                        if checkFiles(js, self.__expected_result)
113                        else "Result is not as expected"
114                    )
115
116    def iteration(self, e):
117        """Iterate over queue to process each bgp element"""
118        # self.redis.xadd()
119
120        if self.verbose:
121            print(json.dumps(e, sort_keys=True, cls=SetEncoder) + "\n", flush=True)
122        if self.json_out:
123            self.json_out.write(
124                json.dumps(e, sort_keys=True, indent=4, cls=SetEncoder) + ",\n"
125            )
126        self.graph.update(e)
127        self.databases.save(e)
128
129    def closeFile(self, file):
130        """Close a JSON file
131
132        Args:
133            file (File): The file to close
134        """
135        if file == sys.stdout or file is None:
136            return
137        file.seek(file.tell() - 1, os.SEEK_SET)
138        file.truncate()
139        file.write("]")
140        file.close()
141
142
143def checkFiles(f1, f2) -> bool:
144    """
145    Check if two json files are equal
146
147    Args:
148        f1, f2 (File): json files
149    """
150
151    f1.seek(0, os.SEEK_SET)
152    f2.seek(0, os.SEEK_SET)
153    json1 = json.load(f1)
154    json2 = json.load(f2)
155    return json1 == json2