FCSC 2022 - MISC - Color Plant 2/2

ColorPlant.py

from pymodbus.client.sync import ModbusTcpClient

UNIT = 1
R = {"to_fill": 32, "valve_id": 0, "regu_id": 32}
G = {"to_fill": 126, "valve_id": 1, "regu_id": 33}
B = {"to_fill": 42, "valve_id": 2, "regu_id": 34}

MAX = 255
M = {"register_id": 6, "max": 100}
F = {"register_id": 10, "max": 200}

class Challenge:
    def __init__(self):
        self.client = ModbusTcpClient('141.94.171.28', port=502)
        self.token = self.getToken()

    def getToken(self):
        return ''.join([chr(c) for c in self.client.read_holding_registers(0, 32, unit=UNIT).registers])

    def readColors(self):
        return self.client.read_input_registers(0, 3, unit=UNIT).registers

    def readCuve(self, cuve):
        return sum(self.client.read_input_registers(cuve, 1, unit=UNIT).registers)
    
    def writeRegu(self, color, val):
        self.client.write_register(color, val, unit=UNIT)
    
    def openColor(self, color):
        self.client.write_coil(color, 1, unit=UNIT)

    def closeColor(self, color):
        self.client.write_coil(color, 0, unit=UNIT)


def main():
    chall = Challenge()
    print('https://color-plant.france-cybersecurity-challenge.fr/'+ chall.token)

    colors = [R, G, B]
    for i in range(len(colors)):
        colors[i]["current"] = colors[i]["close_to_end"] = colors[i]["finished"] = 0

    while 1:
        colors[0]["current"], colors[1]["current"], colors[2]["current"] = chall.readColors()

        if chall.readCuve(F["register_id"]) == F["max"]:
            break
        
        mid = chall.readCuve(M["register_id"])
        if mid == M["max"]:
            chall.openColor(3)
            chall.writeRegu(35, 5)
            while mid != 0:
                mid = chall.readCuve(M["register_id"])

            chall.closeColor(3)

        for i, color in enumerate(colors):
            if mid > 95:
                chall.writeRegu(color["regu_id"], 1)
            
            if not color["finished"]:
                chall.openColor(color["valve_id"])
                if not color["close_to_end"] and mid < 95:
                    chall.writeRegu(color["regu_id"], 5)

            if color["current"] <= MAX - color["to_fill"] + 5:
                chall.writeRegu(color["regu_id"], 1)
                colors[i]["close_to_end"] = 1
                
            if color["current"] <= MAX - color["to_fill"]:
                chall.closeColor(color["valve_id"])
                colors[i]["finished"] = 1

if __name__ == "__main__":
    main()