HiveBrain v1.2.0
Get Started
← Back to all entries
patternpythonModerate

Python bot to answer mathematical questions for a remote server

Submitted by: @import:stackexchange-codereview··
0
Viewed 0 times
answerserverquestionsbotpythonforremotemathematical

Problem

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)

Solution

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

Code Snippets

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))

Context

StackExchange Code Review Q#155384, answer score: 14

Revisions (0)

No revisions yet.