Python bot to answer mathematical questions for a remote server

Question

The MSP Tech Club puts on an annual CTF, available here. I've worked on a solution but was wondering if there's a better way to write it or if I have any particular bad habits showing themselves?

The challenge is to write a script to connect to a server and answer math problems presented to it, it looks like this:

Answering 500 questions in a small enough timeframe reveals this:

I've left the netcat server in the code to allow for testing of changes. Since the CTF is public facing I don't think this poses an issue but feel free to remove if I'm breaking any S/O rules.

#!/usr/bin/python3

import socket
import re

if __name__ == '__main__':
    client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    client.connect(('195.154.53.62', 1337))

    while True:
        data = b''
        while True:
            chunk = client.recv(4096)
            data += chunk
            if len(chunk) < 4096:
                break

        # our flag contains ALEXCTF, once it's revealed print recevied data and exit
        if 'ALEXCTF' in data.decode('utf-8'):
            print(data.decode('utf-8'))
            break

        # \d+ matches a digit (equal to [0-9])
        # .{3} matches any  character, except line terminators exactly three times
        m = re.search('\d+.{3}\d+', data.decode('utf-8'))
        expression = m.group(0)

        #properly handle division
        if '/' in expression:
            expression = expression.replace('/', '//')

        result = eval(expression)

        #print results to screen to see script progress
        print(expression + ' = ' + str(result))

        #encode and transfer
        data = str(result).encode('utf-8') + b'\n'
        client.send(data)

Answer 1

My main concern is the use of eval() without proper validation (well, the regular expression extraction puts some limits onto the possible expression values, but still it can be exploitable). I would either switch to tools like numexpr (there are some other options at this topic), or, taking into account the simplicity of the possible expressions, map the operator strings into actual operators, something like:

import operator


OPERATIONS = {
    '+': operator.add,
    '-': operator.sub,
    '*': operator.mul,
    '/': operator.floordiv  # // operator
}
match = re.search(r'(\d+)\s*([-+*/])\s*(\d+)', data.decode('utf-8'))
if not match:
    raise ValueError("Invalid expression string")

operand1, operation, operand2 = match.groups()

if operation not in OPERATIONS:
    raise ValueError("Invalid operation '%s'" % operation)

result = OPERATIONS[operation](float(operand1), float(operand2))

I'm not sure about converting to floats, may be you need Decimals.

Some other notes:

• it is recommended to use raw strings for regular expression strings
• you can pre-compile the regular expression before going into the loop (even though the Python's regex engine is smart enough to cache the compiled patterns implicitly, see this topic for more information)
• error handling - check if the "match" is not None before calling .group() on it (see how I've done it above)
• fix the typo - "recevied" vs "received"
• comments style - according to PEP8, there has to be a single whitespace after the # character

Answer 2

I'd prefer to see more safety, more structured code, and more abstract symbols instead of magic values. There is some inefficiency due to repeated calls that should be cached.

Safety

You're connecting to the net. You don't set a limit on the size of your input data, and you don't set a limit on the number of times you will respond to the server.

I think it's not unreasonable to abort if the data buffer goes over some upper bound. Realistically, 4k is probably enough. 12k is definitely enough.

Likewise, it's probably reasonable to bail out after 100 iterations through the main loop. No sense burning mips beyond that.

Your call to eval dangerous. Given that you don't check for re match failures, you should do a better job parsing the expression. Please consider defining a set of accepted operations, explicitly handling them, and not using eval. (Which will make your code look like something from Intro to Programming, but that's more appropriate, really.)

Structure

Break your code down into functions. Call them.

Pretty clearly, you "get an expression", "evaluate the expression", and "send the results".

Symbolics

These are obvious:

MAXBUF = 4096
SENTINEL = 'ALEXCTF'
CTF_BOT = ('195.154.53.62', 1337)

There may be more, depending on the expression parsing.

Caching Results

You keep calling data.decode(). Just stick that result in a variable.