Merge pull request #31 from splunk/develop

Develop

Author: weliasz

.circleci/config.yml

@@ -40,6 +40,17 @@ jobs:
           path: build-container
           destination: build-container
 
+  format-code:
+    docker:
+      - image: circleci/python:3.8.6
+    steps:
+      - checkout
+      - run:
+          name: "Check Formatting"
+          command: |
+            pip install black
+            black --check ./ --exclude mibs
+
   release:
     docker:
       - image: circleci/node:12
@@ -68,11 +79,11 @@ jobs:
       - run:
           name: Release Docker Hub
           command: |
-            export VERSION=${CIRCLE_TAG:=CIRCLE_SHA1}      
+            export VERSION=${CIRCLE_TAG:=CIRCLE_SHA1}
             VERSION_DOCKER_MMP=$(echo  $VERSION | sed -n 's/v\([0-9]*\.[0-9]*\.[0-9]*\).*/\1/p')
             VERSION_DOCKER_MM=$(echo  $VERSION | sed -n 's/v\([0-9]*\.[0-9]*\).*/\1/p')
             VERSION_DOCKER_M=$(echo  $VERSION | sed -n 's/v\([0-9]*\).*/\1/p')
-            
+
             curl -sSL https://raw.githubusercontent.com/python-poetry/poetry/master/get-poetry.py | python3 -
             source $HOME/.poetry/env
             poetry build
@@ -112,7 +123,7 @@ jobs:
           command: |
             go get -v -u github.com/tcnksm/ghr
             PATH=$PATH:/usr/local/go/bin
-            export VERSION=${CIRCLE_TAG:=CIRCLE_SHA1}      
+            export VERSION=${CIRCLE_TAG:=CIRCLE_SHA1}
             [[ $VERSION =~ ^v[0-9]+\.[0-9]+\.[0-9]+$ ]] && ISPRE="-prerelease"
             $HOME/go/bin/ghr -t ${GH_TOKEN} -u ${CIRCLE_PROJECT_USERNAME} -r ${CIRCLE_PROJECT_REPONAME} -c ${CIRCLE_SHA1} ${ISPRE} ${VERSION} /tmp/workspace/
       - store_artifacts:
@@ -122,6 +133,8 @@ workflows:
   version: 2
   build_test:
     jobs:
+      - format-code
+
       - build:
           context:
             - gdi-github

.gitignore

@@ -156,3 +156,6 @@ dmypy.json
 .pyre/
 mibs/.DS_Store
 .DS_Store
+
+# Intellij
+.idea

lookups/custom_mib_string_table.csv

@@ -1,6 +1,4 @@
 OID,Label
-1.3.6.1.6.3.1.1.4.1.0,paris
-1.3.6.1.2.1.1.3.0,london
 1.3.6.1.4.1.6141.2.60.37.0.2,splunk-test
 1.3.6.1.4.1.8072.2.3.0.1,mycustomNET-SNMP-EXAMPLES-MIB::netSnmpExampleHeartbeatNotification
 1.3.6.1.4.1.8072.2.3.2.1,mycustomNET-SNMP-EXAMPLES-MIB::netSnmpExampleHeartbeatRate

pyproject.toml

@@ -15,6 +15,7 @@ Flask = "^1.1.2"
 Flask-AutoIndex = "^0.6.6"
 
 [tool.poetry.dev-dependencies]
+mongomock = "^3.22.1"
 
 [tool.poetry.scripts]
 sc4snmp-mib-server = "splunk_connect_for_snmp_mib_server.snmp_mib_server:main"

splunk_connect_for_snmp_mib_server/__init__.py

@@ -1,3 +1,5 @@
 import logging.config
 
-logging.basicConfig(level=logging.INFO, format='%(asctime)s | %(name)s | %(levelname)s | %(message)s')
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s | %(name)s | %(levelname)s | %(message)s"
+)

splunk_connect_for_snmp_mib_server/mib_server.py

@@ -1,58 +1,53 @@
-from flask import Flask, request, abort, jsonify, send_from_directory
+from flask import Flask, request
 from flask_autoindex import AutoIndex
-import os
-import yaml 
-from splunk_connect_for_snmp_mib_server.mongo import MibsRepository, OidsRepository
 from splunk_connect_for_snmp_mib_server.translator import Translator
 import logging
 
 logger = logging.getLogger(__name__)
 
+
 class MibServer:
     def __init__(self, args, server_config):
         self._args = args
         self._server_config = server_config
         self._translator = Translator(server_config)
         self._flask_app = self.build_flask_app()
-    
-    def build_flask_app(self):    
+
+    def build_flask_app(self):
         app = Flask(__name__)
-        # ppath = os.environ['MIBS_SERVER_URL']
-        ppath = self._server_config["snmp"]["mibs"]["mibs_path"]
-        files_index = AutoIndex(app, ppath, add_url_rules=False)
+        mibs_path = self._server_config["snmp"]["mibs"]["mibs_path"]
+        files_index = AutoIndex(app, mibs_path, add_url_rules=False)
 
-        @app.route('/')
+        @app.route("/")
         def hello():
             return "Hello, This is SNMP MIB server!"
 
         # Custom indexing
-        @app.route('/files/')
-        @app.route('/files/<path:path>')
-        def autoindex(path='.'):
+        @app.route("/files/")
+        @app.route("/files/<path:path>")
+        def autoindex(path="."):
             logger.debug(f"path: {path}")
             return files_index.render_autoindex(path)
 
         # Translate oid
-        @app.route('/translation', methods=['POST'])
+        @app.route("/translation", methods=["POST"])
         def translator():
             logger.debug(request.json)
-            var_binds = request.json.get('var_binds')
+            var_binds = request.json.get("var_binds")
             metric = request.args.get("metric")
             logger.debug(f"type of var_binds: {type(var_binds)}")
             logger.debug(f"var_binds: {var_binds}")
             logger.debug(f"type of metric: {str(metric)} -- metric: {metric}")
             if metric == "True":
-                # TODO
-                # If metric is true, var_binds just has one 
+                # if metric is true, var_binds has just one element
                 var_bind = var_binds[0]
-                result = self._translator.format_metric_data(var_bind)                
+                result = self._translator.format_metric_data(var_bind)
             else:
                 result = self._translator.format_trap_event(var_binds)
             return result
-        
+
         return app
-    
-    def run_mib_server(self):
-        # TODO poetry run fails when debug=True
-        self._flask_app.run(host="0.0.0.0",port=self._args.port)
 
+    def run_mib_server(self):
+        # poetry run fails when debug=True
+        self._flask_app.run(host="0.0.0.0", port=self._args.port)

splunk_connect_for_snmp_mib_server/mongo.py

@@ -1,4 +1,4 @@
-from pymongo import MongoClient
+import pymongo
 import os
 import logging
 
@@ -7,10 +7,7 @@
 
 class MibsRepository:
     def __init__(self, mongo_config):
-        """
-        Create a collection in mongodb to store mib files
-        """
-        self._client = MongoClient(
+        self._client = pymongo.MongoClient(
             os.environ["MONGO_SERVICE_SERVICE_HOST"],
             int(os.environ["MONGO_SERVICE_SERVICE_PORT"]),
         )
@@ -23,16 +20,9 @@ def __init__(self, mongo_config):
         self._mibs = self._client[mongo_config["database"]][mongo_config["collection"]]
 
     def upload_files(self, mib_files_dir):
-        """
-        Upload mib files from dir to mongo
-        @param mib_files_dir: the path of the mib files directory
-        """
-        # TODO check duplicate before insert, using filename as PK
         for filename in os.listdir(mib_files_dir):
             file_path = mib_files_dir + "/" + filename
-            # print(file_path)
             with open(file_path, "r") as mib_file:
-                # TODO add try catch, insert only if PK doesn't exit
                 try:
                     self._mibs.insert_one(
                         dict(content=mib_file.read(), filename=filename, _id=filename)
@@ -43,34 +33,22 @@ def upload_files(self, mib_files_dir):
                     )
 
     def search_oid(self, oid):
-        """
-        Search mib file based on oid
-        @param oid: oid, format: "1, 3, 6, 1, 4, 1, 2356, 11, 0, 0, 10000"
-        @return mib filename
-        """
         data = self._mibs.find_one({"content": {"$regex": oid}})
         if data:
             return data["filename"]
         else:
             return None
 
     def delete_mib(self, filename):
-        """
-        Delete mib files based on filename
-        @param filename: mib filename
-        """
         self._mibs.delete_many({"filename": {"$regex": filename}})
 
     def clear(self):
-        """
-        Clear the collection
-        """
         self._mibs.remove()
 
 
 class OidsRepository:
     def __init__(self, mongo_config):
-        self._client = MongoClient(
+        self._client = pymongo.MongoClient(
             os.environ["MONGO_SERVICE_SERVICE_HOST"],
             int(os.environ["MONGO_SERVICE_SERVICE_PORT"]),
         )
@@ -83,7 +61,7 @@ def __init__(self, mongo_config):
         self._oids = self._client[mongo_config["database"]][mongo_config["collection"]]
 
     def contains_oid(self, oid):
-        return self._oids.find({"oid": oid}).count()
+        return self._oids.count_documents({"oid": oid})
 
     def add_oid(self, oid):
         self._oids.insert_one({"oid": oid})

splunk_connect_for_snmp_mib_server/snmp_mib_server.py

@@ -1,7 +1,6 @@
 import argparse
-import os
-import yaml 
-from splunk_connect_for_snmp_mib_server.mongo import MibsRepository, OidsRepository
+import yaml
+from splunk_connect_for_snmp_mib_server.mongo import MibsRepository
 from splunk_connect_for_snmp_mib_server.mib_server import MibServer
 import logging
 
@@ -13,12 +12,13 @@ def upload_mibs(server_config):
     mib_files_dir = server_config["snmp"]["mibs"]["dir"]
     mibs_collection = MibsRepository(mibs_mongo_config)
     # TODO do we need to clean up before loading
-    # Clean up 
+    # Clean up
     mibs_collection.clear()
     # Upload all mib files in specific dir into mongo
     mibs_collection.upload_files(mib_files_dir)
     logger.debug("Uploaded all mib files into mongo!")
 
+
 def main():
     logger.info(f"Startup Config")
     parser = argparse.ArgumentParser()
@@ -35,31 +35,22 @@ def main():
         "--hec_threads", default=10, help="Max http worker thread count", type=int
     )
     parser.add_argument("-c", "--config", default="config.yaml", help="Config File")
-
-
     args = parser.parse_args()
-
     log_level = args.loglevel.upper()
     config_file = args.config
 
     logging.getLogger().setLevel(log_level)
     logger.info(f"Log Level is {log_level}")
     logger.info(f"Config file is {config_file}")
-
     logger.info("Completed Argument parsing")
 
     with open(config_file, "r") as yamlfile:
         server_config = yaml.load(yamlfile, Loader=yaml.FullLoader)
 
-    
-    
-    # Upload mib fiels into mongo
     upload_mibs(server_config)
-       
-    # Init MibServer
     mib_server = MibServer(args, server_config)
     mib_server.run_mib_server()
 
 
 if __name__ == "__main__":
-    main()
+    main()