Merge pull request #532 from NASA-AMMOS/issue-530

Issue #530 - Use SafeLoader for YAML load calls

Author: nttoole

ait/core/bin/ait_bsc.py

@@ -1,5 +1,4 @@
 #!/usr/bin/env python
-
 # Advanced Multi-Mission Operations System (AMMOS) Instrument Toolkit (AIT)
 # Bespoke Link to Instruments and Small Satellites (BLISS)
 #
@@ -13,18 +12,17 @@
 # laws and regulations. User has the responsibility to obtain export licenses,
 # or other export authority as may be required before exporting such
 # information to foreign countries or providing access to foreign persons.
-
 """
 Usage: ait-bsc
 
 Start the ait BSC for capturing network traffic into PCAP files
 and the manager server for RESTful manipulation of active loggers.
 """
-
+import argparse
 import os
 import threading
+
 import yaml
-import argparse
 
 import ait
 from ait.core import bsc
@@ -49,7 +47,7 @@ def main():
 
     else:
         with open(config_file) as log_conf:
-            conf = yaml.load(log_conf, Loader=yaml.Loader)
+            conf = yaml.safe_load(log_conf)
 
         mngr_conf = conf["capture_manager"]
         host = mngr_conf["manager_server"]["host"]

ait/core/cfg.py

@@ -11,25 +11,25 @@
 # laws and regulations. User has the responsibility to obtain export licenses,
 # or other export authority as may be required before exporting such
 # information to foreign countries or providing access to foreign persons.
-
 """
 AIT Configuration
 
 The ait.core.cfg module provides classes and functions to manage
 (re)configurable aspects of AIT via a YAML configuration file.
 
 """
-
 import os
 import platform
+import re
 import sys
 import time
-import re
-import yaml
 from io import IOBase
 
+import yaml
+
 import ait
-from ait.core import log, util
+from ait.core import log
+from ait.core import util
 
 
 DEFAULT_PATH_VARS = {
@@ -61,8 +61,9 @@ def expand_config_paths(
 
             for p in cleaned:
                 if not os.path.exists(p):
-                    msg = "Config parameter {}.{} specifies nonexistent path {}".format(
-                        parameter_key, name, p
+                    msg = (
+                        "Config parameter {}.{} specifies nonexistent "
+                        "path {}".format(parameter_key, name, p)
                     )
                     log.warn(msg)
 
@@ -176,7 +177,7 @@ def load_yaml(filename=None, data=None):
         if filename:
             data = open(filename, "rt")
 
-        config = yaml.load(data, Loader=yaml.Loader)
+        config = yaml.safe_load(data)
 
         if isinstance(data, IOBase):
             data.close()
@@ -204,13 +205,13 @@ class AitConfigError(Exception):
     pass
 
 
-class AitConfigMissing(Exception):
+class AitConfigMissingError(Exception):
     """Raised when a AIT configuration parameter is missing."""
 
     def __init__(self, param):
         values = param, ait.config._filename
         format = "The parameter %s is missing from config.yaml (%s)."
-        super(AitConfigMissing, self).__init__(format % values)
+        super(AitConfigMissingError, self).__init__(format % values)
         self.param = param
 
 
@@ -342,7 +343,7 @@ def _datapaths(self):
                 paths["mib"] = data["mib"]["path"]
 
         except KeyError as e:
-            raise AitConfigMissing(str(e))
+            raise AitConfigMissingError(str(e))
         except Exception as e:
             raise AitConfigError("Error reading data paths: %s" % e)
 

ait/core/cmd.py

@@ -11,22 +11,23 @@
 # laws and regulations. User has the responsibility to obtain export licenses,
 # or other export authority as may be required before exporting such
 # information to foreign countries or providing access to foreign persons.
-
 """
 AIT Commands
 
 The ait.core.cmd module provides commands and command dictionaries.
 Dictionaries contain command and argument definitions.
 """
-
 import os
-import pkg_resources
 import struct
-import yaml
 from io import IOBase
 
+import pkg_resources
+import yaml
+
 import ait
-from ait.core import json, util, log
+from ait.core import json
+from ait.core import log
+from ait.core import util
 
 
 MAX_CMD_WORDS = 54
@@ -124,7 +125,11 @@ def encode(self, value):
 
         if type(value) == str and self.enum and value in self.enum:
             value = self.enum[value]
-        return self.type.encode(*value) if type(value) in [tuple, list] else self.type.encode(value)
+        return (
+            self.type.encode(*value)
+            if type(value) in [tuple, list]
+            else self.type.encode(value)
+        )
 
     def slice(self, offset=0):
         """Returns a Python slice object (e.g. for array indexing) indicating
@@ -486,7 +491,7 @@ def load(self, content):
             else:
                 stream = content
 
-            cmds = yaml.load(stream, Loader=yaml.Loader)
+            cmds = yaml.safe_load(stream)
             cmds = handle_includes(cmds)
             for cmd in cmds:
                 self.add(cmd)
@@ -503,7 +508,7 @@ def getDefaultCmdDict(reload=False):  # noqa
 
 
 def getDefaultDict(reload=False):  # noqa
-    create_cmd_dict_func = globals().get('createCmdDict', None)
+    create_cmd_dict_func = globals().get("createCmdDict", None)
     loader = create_cmd_dict_func if create_cmd_dict_func else CmdDict
     return util.getDefaultDict(__name__, "cmddict", loader, reload)
 
@@ -559,13 +564,13 @@ def YAMLCtor_include(loader, node):  # noqa
     name = os.path.join(os.path.dirname(loader.name), node.value)
     data = None
     with open(name, "r") as f:
-        data = yaml.load(f)
+        data = yaml.safe_load(f)
     return data
 
 
-yaml.add_constructor("!include", YAMLCtor_include)
-yaml.add_constructor("!Command", YAMLCtor_CmdDefn)
-yaml.add_constructor("!Argument", YAMLCtor_ArgDefn)
-yaml.add_constructor("!Fixed", YAMLCtor_ArgDefn)
+yaml.SafeLoader.add_constructor("!include", YAMLCtor_include)
+yaml.SafeLoader.add_constructor("!Command", YAMLCtor_CmdDefn)
+yaml.SafeLoader.add_constructor("!Argument", YAMLCtor_ArgDefn)
+yaml.SafeLoader.add_constructor("!Fixed", YAMLCtor_ArgDefn)
 
 util.__init_extensions__(__name__, globals())

ait/core/evr.py

@@ -11,20 +11,22 @@
 # laws and regulations. User has the responsibility to obtain export licenses,
 # or other export authority as may be required before exporting such
 # information to foreign countries or providing access to foreign persons.
-
 """
 AIT Event Record (EVR) Reader
 
 The ait.core.evr module is used to read the EVRs from a YAML file.
 """
-
 import os
-import pkg_resources
 import re
+
+import pkg_resources
 import yaml
 
 import ait.core
-from ait.core import dtype, json, log, util
+from ait.core import dtype
+from ait.core import json
+from ait.core import log
+from ait.core import util
 
 
 class EVRDict(dict):
@@ -59,7 +61,7 @@ def load(self, content):
             stream = content
 
         try:
-            evrs = yaml.load(stream, Loader=yaml.Loader)
+            evrs = yaml.safe_load(stream)
         except IOError as e:
             msg = "Could not load EVR YAML '{}': '{}'".format(stream, str(e))
             log.error(msg)
@@ -196,7 +198,9 @@ def format_message(self, evr_hist_data):
                 data_chunks.append(d)
             # TODO: Make this not suck
             except Exception:
-                msg = "Unable to format EVR Message with data {}".format(evr_hist_data)
+                msg = "Unable to format EVR Message with data " "{}".format(
+                    evr_hist_data
+                )
                 log.error(msg)
                 raise ValueError(msg)
 
@@ -212,9 +216,9 @@ def format_message(self, evr_hist_data):
         if len(formatters) == 0:
             return self._message
         else:
-            # Python format strings cannot handle size formatter information. So something
-            # such as %llu needs to be adjusted to be a valid identifier in python by
-            # removing the size formatter.
+            # Python format strings cannot handle size formatter information.
+            # So something such as %llu needs to be adjusted to be a valid
+            # identifier in python by removing the size formatter.
             msg = self._message
             for f in formatters:
                 if len(f) > 1:
@@ -237,6 +241,6 @@ def YAMLCtor_EVRDefn(loader, node):  # noqa
     return createEVRDefn(**fields)  # noqa
 
 
-yaml.add_constructor("!EVR", YAMLCtor_EVRDefn)
+yaml.SafeLoader.add_constructor("!EVR", YAMLCtor_EVRDefn)
 
 util.__init_extensions__(__name__, globals())

ait/core/limits.py

@@ -11,7 +11,6 @@
 # laws and regulations. User has the responsibility to obtain export licenses,
 # or other export authority as may be required before exporting such
 # information to foreign countries or providing access to foreign persons.
-
 """
 AIT Limits
 
@@ -20,20 +19,26 @@
 The expected limits.yaml should follow this schema:
 
 - !Limit
-  source:      -- telemetry source for the limit. should follow format 'Packet.field_name'
+  source:      -- telemetry source for the limit. should follow format
+                  'Packet.field_name'
   desc:        -- description of the limit
-  units:       -- the units used for possible conversion depending on the units set in the
-                  telemetry dictionary
+  units:       -- the units used for possible conversion depending on
+                  the units set in the telemetry dictionary
   lower:       -- lower limits
-    error:     -- trigger error if telemetry value exceeds this lower bound (exclusive)
-    warn:      -- trigger warning if telemetry value exceeds this lower bound (exclusive)
+    error:     -- trigger error if telemetry value exceeds this lower
+                  bound (exclusive)
+    warn:      -- trigger warning if telemetry value exceeds this lower
+                  bound (exclusive)
   upper:       -- upper limits
-    error:     -- trigger error if telemetry value exceeds this upper bound (exclusive)
-    warn:      -- trigger warning if telemetry value exceeds this upper bound (exclusive)
+    error:     -- trigger error if telemetry value exceeds this upper
+                  bound (exclusive)
+    warn:      -- trigger warning if telemetry value exceeds this upper
+                  bound (exclusive)
   value:       -- enumerated values to trigger error/warning
     error:     -- trigger error if telemetry value == or in list of strings
     warn:      -- trigger warning if telemetry value == or in list of strings
-  when:        -- when condition for specifying the necessary state when this limit applies
+  when:        -- when condition for specifying the necessary state when this
+                  limit applies
   persist:     -- number of seconds the value must persist before limits trigger
 
 For example:
@@ -62,14 +67,15 @@
         - BAR
 
 """
-
 import os
+from io import IOBase
+
 import pkg_resources
 import yaml
-from io import IOBase
 
 import ait
-from ait.core import json, util
+from ait.core import json
+from ait.core import util
 
 
 class Thresholds(json.SlotSerializer, object):
@@ -201,7 +207,7 @@ def load(self, content):
             else:
                 stream = content
 
-            limits = yaml.load(stream, Loader=yaml.Loader)
+            limits = yaml.safe_load(stream)
 
             for lmt in limits:
                 self.add(lmt)
@@ -230,6 +236,6 @@ def YAMLCtor_LimitDefinition(loader, node):  # noqa
     return createLimitDefinition(**fields)  # noqa
 
 
-yaml.add_constructor("!Limit", YAMLCtor_LimitDefinition)
+yaml.SafeLoader.add_constructor("!Limit", YAMLCtor_LimitDefinition)
 
 util.__init_extensions__(__name__, globals())