#!/usr/bin/env python
 #
 # Licensed under the Apache License, Version 2.0 (the "License"); you may
 # not use this file except in compliance with the License. You may obtain
 # a copy of the License at
 #
 #      http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 # License for the specific language governing permissions and limitations
 # under the License.
 
 # This small script updates the libvirt CPU map to add a gate64 cpu model
 # that can be used to enable a common 64bit capable feature set across
 # devstack nodes so that features like nova live migration work.
 
 import sys
 import xml.etree.ElementTree as ET
 from xml.dom import minidom
 
 
 def update_cpu_map(tree):
     root = tree.getroot()
     cpus = root#.find("cpus")
     x86 = None
     for arch in cpus.findall("arch"):
         if arch.get("name") == "x86":
             x86 = arch
             break
     if x86 is not None:

         # Create a gate64 cpu model that is core2duo less monitor, pse36,
         # vme, and ssse3.

         gate64 = ET.SubElement(x86, "model")
         gate64.set("name", "gate64")
         ET.SubElement(gate64, "vendor").set("name", "Intel")
         ET.SubElement(gate64, "feature").set("name", "fpu")
         ET.SubElement(gate64, "feature").set("name", "de")
         ET.SubElement(gate64, "feature").set("name", "pse")
         ET.SubElement(gate64, "feature").set("name", "tsc")
         ET.SubElement(gate64, "feature").set("name", "msr")
         ET.SubElement(gate64, "feature").set("name", "pae")
         ET.SubElement(gate64, "feature").set("name", "mce")
         ET.SubElement(gate64, "feature").set("name", "cx8")
         ET.SubElement(gate64, "feature").set("name", "apic")
         ET.SubElement(gate64, "feature").set("name", "sep")
         ET.SubElement(gate64, "feature").set("name", "pge")
         ET.SubElement(gate64, "feature").set("name", "cmov")
         ET.SubElement(gate64, "feature").set("name", "pat")
         ET.SubElement(gate64, "feature").set("name", "mmx")
         ET.SubElement(gate64, "feature").set("name", "fxsr")
         ET.SubElement(gate64, "feature").set("name", "sse")
         ET.SubElement(gate64, "feature").set("name", "sse2")
         ET.SubElement(gate64, "feature").set("name", "mtrr")
         ET.SubElement(gate64, "feature").set("name", "mca")
         ET.SubElement(gate64, "feature").set("name", "clflush")
         ET.SubElement(gate64, "feature").set("name", "pni")
         ET.SubElement(gate64, "feature").set("name", "nx")
         ET.SubElement(gate64, "feature").set("name", "syscall")
         ET.SubElement(gate64, "feature").set("name", "lm")
 
 
 def format_xml(root):
     # Adapted from http://pymotw.com/2/xml/etree/ElementTree/create.html
     # thank you dhellmann
     rough_string = ET.tostring(root, encoding="UTF-8")
     dom_parsed = minidom.parseString(rough_string)
     return dom_parsed.toprettyxml("  ", encoding="UTF-8")
 
 
 def main():
     if len(sys.argv) != 2:
         raise Exception("Must pass path to cpu_map.xml to update")
     cpu_map = sys.argv[1]
     tree = ET.parse(cpu_map)
     for model in tree.getroot().iter("model"):
         if model.get("name") == "gate64":
             # gate64 model is already present
             return
     update_cpu_map(tree)
     pretty_xml = format_xml(tree.getroot())
     with open(cpu_map, 'w') as f:
         f.write(pretty_xml)
 
 
 if __name__ == "__main__":
     main()