[pyyaml] allow patching of unsafe pyyaml operations

[utils] changing yaml to ddyaml to avoid clash

[pyyaml] allow reversing of monkey patch

.rubocop.yml

@@ -24,12 +24,14 @@ Style/Documentation:
 Style/SingleLineBlockParams:
   Enabled: false
 
+Style/FrozenStringLiteralComment:
+  Enabled: false
+
 Lint/RescueWithoutErrorClass:
   Enabled: false
 
 BlockLength:
   Max: 110
 
-Style/EndOfLine:
+Layout/EndOfLine:
   Enabled: false
-

agent.py

@@ -63,6 +63,7 @@
 from utils.service_discovery.sd_backend import get_sd_backend
 from utils.watchdog import Watchdog
 from utils.windows_configuration import get_sdk_integration_paths
+from utils.ddyaml import monkey_patch_pyyaml
 
 # Constants
 PID_NAME = "dd-agent"
@@ -502,6 +503,10 @@ def main():
     hostname = get_hostname(agentConfig)
     in_developer_mode = agentConfig.get('developer_mode')
 
+    # do this early on
+    if agentConfig.get('disable_unsafe_yaml'):
+        monkey_patch_pyyaml()
+
     COMMANDS_AGENT = [
         'start',
         'stop',

checks/__init__.py

@@ -30,11 +30,12 @@
 # project
 from checks import check_status
 from config import AGENT_VERSION, _is_affirmative
-from util import get_next_id, yLoader
+from util import get_next_id
 from utils.hostname import get_hostname
 from utils.proxy import get_proxy
 from utils.profile import pretty_statistics
 from utils.proxy import get_no_proxy_from_env, config_proxy_skip
+from utils.ddyaml import yLoader
 
 
 log = logging.getLogger(__name__)

config.py

@@ -595,6 +595,10 @@ def get_config(parse_args=True, cfg_path=None, options=None, can_query_registry=
         if config.has_option("Main", "skip_ssl_validation"):
             agentConfig["skip_ssl_validation"] = _is_affirmative(config.get("Main", "skip_ssl_validation"))
 
+        agentConfig["disable_unsafe_yaml"] = True
+        if config.has_option("Main", "disable_unsafe_yaml"):
+            agentConfig["disable_unsafe_yaml"] = _is_affirmative(config.get("Main", "disable_unsafe_yaml"))
+
         agentConfig["collect_instance_metadata"] = True
         if config.has_option("Main", "collect_instance_metadata"):
             agentConfig["collect_instance_metadata"] = _is_affirmative(config.get("Main", "collect_instance_metadata"))

jmxfetch.py

@@ -28,9 +28,9 @@
     _is_affirmative,
 )
 from daemon import ProcessRunner
-from util import yLoader
 from utils.jmx import JMX_FETCH_JAR_NAME, JMXFiles
 from utils.platform import Platform
+from utils.ddyaml import monkey_patch_pyyaml, yLoader
 
 log = logging.getLogger('jmxfetch')
 
@@ -488,6 +488,10 @@ def main(config_path=None):
     """ JMXFetch main entry point """
     confd_path, agent_config = init(config_path)
 
+    # do this early on
+    if agent_config.get('disable_unsafe_yaml'):
+        monkey_patch_pyyaml()
+
     jmx = JMXFetch(confd_path, agent_config)
     return jmx.run()
 

tests/core/test_utils_yaml.py

@@ -0,0 +1,94 @@
+# stdlib
+import os
+import unittest
+import tempfile
+
+# project
+import yaml
+
+from utils.ddyaml import (
+    monkey_patch_pyyaml,
+    monkey_patch_pyyaml_reverse,
+    safe_yaml_dump_all,
+    safe_yaml_load_all,
+    safe_yaml_load,
+    yDumper,
+)
+
+FIXTURE_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'fixtures', 'checks')
+
+class Dummy(object):
+    def __init__(self):
+        self.foo = 1
+        self.bar = 'a'
+        self.qux = {self.foo: self.bar}
+
+    def get_foo(self):
+        return self.foo
+
+    def get_bar(self):
+        return self.bar
+
+    def get_qux(self):
+        return self.qux
+
+
+class UtilsYAMLTest(unittest.TestCase):
+
+    def setUp(self):
+        monkey_patch_pyyaml()
+
+    def tearDown(self):
+        monkey_patch_pyyaml_reverse()
+
+    def test_monkey_patch(self):
+        self.assertTrue(yaml.dump_all == safe_yaml_dump_all)
+        self.assertTrue(yaml.load_all == safe_yaml_load_all)
+        self.assertTrue(yaml.load == safe_yaml_load)
+
+    def test_load(self):
+        conf = os.path.join(FIXTURE_PATH, "valid_conf.yaml")
+        with open(conf) as f:
+            stream = f.read()
+
+            yaml_config_safe = safe_yaml_load(stream)
+            yaml_config_native = yaml.load(stream)
+            self.assertTrue(yaml_config_safe is not None)
+            self.assertTrue(yaml_config_native is not None)
+            self.assertTrue(yaml_config_native == yaml_config_safe)
+
+            yaml_config_safe = [entry for entry in safe_yaml_load_all(stream)]
+            yaml_config_native = [entry for entry in yaml.load_all(stream)]
+            self.assertTrue(yaml_config_safe is not [])
+            self.assertTrue(yaml_config_native is not [])
+            self.assertTrue(len(yaml_config_safe) == len(yaml_config_native))
+            for safe, native in zip(yaml_config_safe, yaml_config_native):
+                self.assertTrue(safe == native)
+
+    def test_unsafe(self):
+        dummy = Dummy()
+
+        with self.assertRaises(yaml.representer.RepresenterError):
+            yaml.dump_all([dummy])
+
+        with self.assertRaises(yaml.representer.RepresenterError):
+            yaml.dump(dummy, Dumper=yDumper)
+
+        # reverse monkey patch and try again
+        monkey_patch_pyyaml_reverse()
+
+        with tempfile.TemporaryFile(suffix='.yaml') as f:
+            yaml.dump_all([dummy], stream=f)
+            f.seek(0)  # rewind
+
+            doc_unsafe = yaml.load(f)
+            self.assertTrue(type(doc_unsafe) is Dummy)
+
+            monkey_patch_pyyaml()
+            with self.assertRaises(yaml.constructor.ConstructorError):
+                f.seek(0)  # rewind
+                safe_yaml_load(f)
+
+            with self.assertRaises(yaml.constructor.ConstructorError):
+                f.seek(0)  # rewind
+                yaml.load(f)

util.py

@@ -10,13 +10,6 @@
 
 # 3p
 import yaml  # noqa, let's guess, probably imported somewhere
-try:
-    from yaml import CSafeLoader as yLoader
-    from yaml import CSafeDumper as yDumper
-except ImportError:
-    # On source install C Extensions might have not been built
-    from yaml import SafeLoader as yLoader  # noqa, imported from here elsewhere
-    from yaml import SafeDumper as yDumper  # noqa, imported from here elsewhere
 
 # These classes are now in utils/, they are just here for compatibility reasons,
 # if a user actually uses them in a custom check
@@ -26,6 +19,7 @@
 from utils.platform import Platform, get_os # noqa, see ^^^
 from utils.proxy import get_proxy # noqa, see ^^^
 from utils.timer import Timer # noqa, see ^^^
+from utils.ddyaml import yLoader
 
 COLON_NON_WIN_PATH = re.compile(':(?!\\\\)')
 
@@ -104,6 +98,7 @@ def get_next_id(name):
 class NoInstancesFound(Exception):
     pass
 
+
 def check_yaml(conf_path):
     with open(conf_path) as f:
         check_config = yaml.load(f.read(), Loader=yLoader)

utils/ddyaml.py

@@ -0,0 +1,101 @@
+# (C) Datadog, Inc. 2011-2016
+# All rights reserved
+# Licensed under Simplified BSD License (see LICENSE)
+
+# stdlib
+import logging
+
+# 3p
+import yaml  # noqa, let's guess, probably imported somewhere
+try:
+    from yaml import CSafeLoader as yLoader
+    from yaml import CSafeDumper as yDumper
+except ImportError:
+    # On source install C Extensions might have not been built
+    from yaml import SafeLoader as yLoader  # noqa, imported from here elsewhere
+    from yaml import SafeDumper as yDumper  # noqa, imported from here elsewhere
+
+log = logging.getLogger(__name__)
+
+pyyaml_load = None
+pyyaml_load_all = None
+pyyaml_dump_all = None
+
+
+def safe_yaml_dump_all(documents, stream=None, Dumper=yDumper,
+        default_style=None, default_flow_style=None,
+        canonical=None, indent=None, width=None,
+        allow_unicode=None, line_break=None,
+        encoding='utf-8', explicit_start=None, explicit_end=None,
+        version=None, tags=None):
+    if Dumper != yDumper:
+        log.warning("Unsafe dumping of YAML has been disabled - using safe dumper instead")
+
+    if pyyaml_dump_all:
+        return pyyaml_dump_all(documents, stream, yDumper,
+            default_style, default_flow_style,
+            canonical, indent, width,
+            allow_unicode, line_break,
+            encoding, explicit_start, explicit_end,
+            version, tags)
+
+    return yaml.dump_all(documents, stream, yDumper,
+        default_style, default_flow_style,
+        canonical, indent, width,
+        allow_unicode, line_break,
+        encoding, explicit_start, explicit_end,
+        version, tags)
+
+def safe_yaml_load(stream, Loader=yLoader):
+    if Loader != yLoader:
+        log.warning("Unsafe loading of YAML has been disabled - using safe loader instead")
+
+    if pyyaml_load:
+        return pyyaml_load(stream, Loader=yLoader)
+
+    return yaml.load(stream, Loader=yLoader)
+
+def safe_yaml_load_all(stream, Loader=yLoader):
+    if Loader != yLoader:
+        log.warning("Unsafe loading of YAML has been disabled - using safe loader instead")
+
+    if pyyaml_load_all:
+        return pyyaml_load_all(stream, Loader=yLoader)
+
+    return yaml.load_all(stream, Loader=yLoader)
+
+def monkey_patch_pyyaml():
+    global pyyaml_load
+    global pyyaml_load_all
+    global pyyaml_dump_all
+
+    if not pyyaml_load:
+        log.info("monkey patching yaml.load...")
+        pyyaml_load = yaml.load
+        yaml.load = safe_yaml_load
+    if not pyyaml_load_all:
+        log.info("monkey patching yaml.load_all...")
+        pyyaml_load_all = yaml.load_all
+        yaml.load_all = safe_yaml_load_all
+    if not pyyaml_dump_all:
+        log.info("monkey patching yaml.dump_all... (affects all yaml dump operations)")
+        pyyaml_dump_all = yaml.dump_all
+        yaml.dump_all = safe_yaml_dump_all
+
+def monkey_patch_pyyaml_reverse():
+    global pyyaml_load
+    global pyyaml_load_all
+    global pyyaml_dump_all
+
+    if pyyaml_load:
+        log.info("reversing monkey patch for yaml.load...")
+        yaml.load = pyyaml_load
+        pyyaml_load = None
+    if pyyaml_load_all:
+        log.info("reversing monkey patch for yaml.load_all...")
+        yaml.load_all = pyyaml_load_all
+        pyyaml_load_all = None
+    if pyyaml_dump_all:
+        log.info("reversing monkey patch for yaml.dump_all... (affects all yaml dump operations)")
+        yaml.dump_all = pyyaml_dump_all
+        pyyaml_dump_all = None

utils/jmx.py

@@ -13,9 +13,9 @@
 
 # datadog
 from config import _windows_commondata_path, get_confd_path, JMX_VERSION
-from util import yDumper
 from utils.pidfile import PidFile
 from utils.platform import Platform
+from utils.ddyaml import yDumper
 
 # JMXFetch java version
 JMX_FETCH_JAR_NAME = "jmxfetch-{ver}-jar-with-dependencies.jar".format(ver=JMX_VERSION)

win32/gui.py

@@ -84,9 +84,9 @@
     get_logging_config,
     get_version
 )
-from util import yLoader
 from utils.flare import Flare
 from utils.platform import Platform
+from utils.ddyaml import yLoader
 
 # Constants describing the agent state
 AGENT_RUNNING = 0