#!/usr/bin/python
#
# Script to add Chameleon GenerateCStates and GeneratePStates
#
# Copyright (C) tonymacx86 LLC

import sys
import os

vol = str(sys.argv[1])
boot = "/Extra/org.chameleon.Boot.plist"
plist = vol + boot
infile = open(plist, "r")

CStatesCheck = False
PStatesCheck = False

body = ""

for line in infile:
    if "<key>GenerateCStates</key>" in line:
        CStatesCheck = True

    if "<key>GeneratePStates</key>" in line:
        PStatesCheck = True
        
    # if we finish the tags and haven't seen either GenerateStates
    if "</dict>" in line:
        line = ""
        if CStatesCheck == False:
            line += "        <key>GenerateCStates</key>\n"
            line += "        <string>Yes</string>\n"
            CStatesCheck = True

        if PStatesCheck == False:
            line += "        <key>GeneratePStates</key>\n"
            line += "        <string>Yes</string>\n"
            PStatesCheck = True
            
        line += "</dict>\n"
        
    body += line
    
infile.close()

outfile = open(plist, "w")
outfile.write(body)
outfile.close()