diff --git a/model-optimizer/mo/back/ie_ir_ver_2/emitter.py b/model-optimizer/mo/back/ie_ir_ver_2/emitter.py index 47b279034d3110..ec9b737e464f8f 100644 --- a/model-optimizer/mo/back/ie_ir_ver_2/emitter.py +++ b/model-optimizer/mo/back/ie_ir_ver_2/emitter.py @@ -2,8 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 import hashlib -from xml.etree.ElementTree import Element, SubElement, tostring, ElementTree +from defusedxml import defuse_stdlib +import defusedxml.ElementTree as ET from defusedxml.minidom import parseString from mo.graph.graph import * @@ -12,6 +13,13 @@ from mo.utils.utils import refer_to_faq_msg from mo.utils.version import get_version +# defuse_stdlib provide patched version of xml.etree.ElementTree which allows to use objects from xml.etree.ElementTree +# in a safe manner without including unsafe xml.etree.ElementTree +ET_defused = defuse_stdlib()[ET] +Element = ET_defused.Element +SubElement = ET_defused.SubElement +tostring = ET_defused.tostring + def serialize_constants(graph: Graph, bin_file_name: str, data_type=np.float32): """ @@ -444,8 +452,7 @@ def append_ir_info(file: str, meta_info: dict = dict(), mean_data: [list, None] path_to_xml = file + ".xml" path_to_bin = file + ".bin" - et = ElementTree() - et.parse(path_to_xml) + et = ET.parse(path_to_xml) net = et.getroot() if mean_data: @@ -462,4 +469,4 @@ def append_ir_info(file: str, meta_info: dict = dict(), mean_data: [list, None] pretty_xml_as_string = parseString(tostring(net)).toprettyxml() with open(path_to_xml, 'wb') as file: - file.write(bytes(pretty_xml_as_string, "UTF-8")) \ No newline at end of file + file.write(bytes(pretty_xml_as_string, "UTF-8")) diff --git a/model-optimizer/mo/middle/passes/tensor_names.py b/model-optimizer/mo/middle/passes/tensor_names.py index 4e291023679ff5..d8bb5591e0c783 100644 --- a/model-optimizer/mo/middle/passes/tensor_names.py +++ b/model-optimizer/mo/middle/passes/tensor_names.py @@ -1,12 +1,19 @@ # Copyright (C) 2018-2021 Intel Corporation # SPDX-License-Identifier: Apache-2.0 -from xml.etree.ElementTree import Element, SubElement, tostring - +from defusedxml import defuse_stdlib from defusedxml.minidom import parseString +import defusedxml.ElementTree as ET from mo.graph.graph import Node, Graph +# defuse_stdlib provide patched version of xml.etree.ElementTree which allows to use objects from xml.etree.ElementTree +# in a safe manner without including unsafe xml.etree.ElementTree +ET_defused = defuse_stdlib()[ET] +Element = ET_defused.Element +SubElement = ET_defused.SubElement +tostring = ET_defused.tostring + def propagate_op_name_to_tensor(graph: Graph): for node in graph.nodes(): diff --git a/model-optimizer/mo/utils/ir_engine/ir_engine.py b/model-optimizer/mo/utils/ir_engine/ir_engine.py index b124e6d30046a6..35d95486c5f734 100644 --- a/model-optimizer/mo/utils/ir_engine/ir_engine.py +++ b/model-optimizer/mo/utils/ir_engine/ir_engine.py @@ -5,7 +5,9 @@ import logging as log import os import sys -import xml.etree.ElementTree as ET + +from defusedxml import defuse_stdlib +import defusedxml.ElementTree as ET from argparse import Namespace from collections import namedtuple, defaultdict from pathlib import Path @@ -17,6 +19,9 @@ log.basicConfig(format="[ %(levelname)s ] %(message)s", level=log.DEBUG, stream=sys.stdout) +# defuse_stdlib provide patched version of xml.etree.ElementTree which allows to use objects from xml.etree.ElementTree +# in a safe manner without including unsafe xml.etree.ElementTree +ElementTree = defuse_stdlib()[ET].ElementTree class IREngine(object): def __init__(self, path_to_xml: str, path_to_bin=None, precision="FP32", xml_tree=None): @@ -88,7 +93,6 @@ def __load_xml(self): elif elem.tag in ['version', 'cli_params']: self.meta_data['quantization_parameters'][elem.tag] = elem.attrib['value'] - self.graph.graph['cmd_params'] = Namespace(**self.meta_data) # TODO check what we need all this attrs if len(statistics): @@ -237,7 +241,7 @@ def __load_layer(self, layer): body_ir = IREngine(path_to_xml=None, path_to_bin=self.path_to_bin, - xml_tree=ET.ElementTree(xml_body_child[0])) + xml_tree=ElementTree(xml_body_child[0])) self.graph.graph['hashes'].update(body_ir.graph.graph['hashes']) # Find port_map section and take an input_port_map & output_port_map diff --git a/model-optimizer/unit_tests/mo/utils/ir_reader/restore_graph_test.py b/model-optimizer/unit_tests/mo/utils/ir_reader/restore_graph_test.py new file mode 100644 index 00000000000000..988014f3840ef5 --- /dev/null +++ b/model-optimizer/unit_tests/mo/utils/ir_reader/restore_graph_test.py @@ -0,0 +1,35 @@ +# Copyright (C) 2018-2021 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +import os +import unittest +import tempfile + +from mo.utils.ir_reader.restore_graph import restore_graph_from_ir +from defusedxml.common import EntitiesForbidden + + +class TestIRReader(unittest.TestCase): + def setUp(self): + self.xml_bomb = b'\n' \ + b'\n' \ + b' \n' \ + b' \n' \ + b' \n' \ + b' \n' \ + b' \n' \ + b' \n' \ + b' \n' \ + b' \n' \ + b' \n' \ + b' \n' \ + b']>\n' \ + b'&lol9;' + + def test_read_xml_bomb(self): + bomb_file = tempfile.NamedTemporaryFile(delete=False) + bomb_file.write(self.xml_bomb) + bomb_file.close() + self.assertRaises(EntitiesForbidden, restore_graph_from_ir, bomb_file.name) + os.remove(bomb_file.name)