#!/usr/bin/env python2
"""A basic exploit example.

Input:
    An IP address as first argument:
Output:
    A string composed of the service name, target IP
    address, flag and a timestamp in the following format:
    $service|$ip|$flag|$timestamp"""
import sys
import time
import re
from twisted.internet import reactor, protocol


if len(sys.argv) < 2:
    print('Please provide an IP/Host as argument.')
    sys.exit(1)


SERVICE_NAME = 'smartgrid'
SERVICE_ADDR = sys.argv[1]
SERVICE_PORT = int(sys.argv[2])

FLAG_REGEX = r"(\w{31}=)"
FLAG_GREP = re.compile(FLAG_REGEX)


def outputFlag(flag):
    print('{}|{}|{}|{}'.format(SERVICE_NAME, SERVICE_ADDR, flag.strip(), int(time.time())))


class ServiceExploiterProtocol(protocol.Protocol):

    def __init__(self):
        self.flags = None

    def connectionMade(self):
        print('Connected to service...')

    def dataReceived(self, incoming):
        self.flags = FLAG_GREP.findall(incoming)
        if len(self.flags):
            self.transport.loseConnection()
            return

        self.writeLn('magic data')

    def connectionLost(self, reason):
        if len(self.flags):
            for flag in self.flags:
                outputFlag(flag)

    def writeLn(self, msg):
        self.transport.write('{}\n'.format(msg))


class ServiceExploiterFactory(protocol.Factory):
    protocol = ServiceExploiterProtocol

    def startedConnecting(self, connector):
        pass

    def buildProtocol(self, addr):
        return ServiceExploiterProtocol()

    def clientConnectionLost(self, connector, reason):
        pass

    def clientConnectionFailed(self, connector, reason):
        print('Connection failed: {}'.format(reason))


if __name__ == '__main__':
    try:
        print('Starting exploiter...')
        reactor.connectTCP(SERVICE_ADDR, SERVICE_PORT, ServiceExploiterFactory())
        reactor.run()
        sys.exit(0)
    except KeyboardInterrupt:
        reactor.stop()
        sys.exit(0)
    except Exception as e:
        print(e)
        sys.exit(1)
