python-emwin/emwin.py at master · jhavard/python-emwin · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
from time import strptime, mktime, time
import logging
import sys


handler = logging.StreamHandler()
handler.setFormatter(logging.Formatter('%(asctime)s %(name)s %(levelname)s %(message)s'))
log = logging.getLogger('emwin')
log.addHandler(handler)
log.setLevel(logging.DEBUG)


class Connection(object):
    def __init__(self, sock):
        self.sock = sock
        self.ident = 'ByteBlast Client|NM-emwin@synack.me|V1'
        self.ident = ''.join([chr(ord(x) ^ 0xFF) for x in self.ident])

    def __iter__(self):
        buf = ''
        last_ident = 0
        while True:
            now = int(time())
            if (now - last_ident) > 300:
                log.info('Sending ident packet')
                last_ident = now
                self.sock.sendall(self.ident)

            buf += self.sock.recv(1116)
            if buf == '':
                break
            while len(buf) >= 1116:
                if not buf.startswith('\xFF\xFF\xFF\xFF\xFF\xFF'):
                    offset = buf.find('\xFF\xFF\xFF\xFF\xFF\xFF')
                    if offset == -1:
                        log.info('Sync marker missing! Abort!')
                        break
                        buf = ''
                    buf = buf[offset:]
                    log.info('Discarding %i bytes before sync marker' % offset)

                try:
                    packet = Packet(buf[:1116])
                    log.debug(str(packet))
                    yield packet
                except:
                    log.error(sys.exc_info()[1])
                    break
                buf = buf[1116:]
        log.error('Connection closed by remote host')
        self.sock.close()


class Packet(object):
    def __init__(self, data):
        self.data = data
        self.parse()

    def parse(self):
        self.data = ''.join([chr(ord(x) ^ 0xFF) for x in self.data])
        self.header = self.parse_header(self.data[:86])
        self.filename = self.header['PF']
        self.block = int(self.header['PN'])
        self.total_blocks = int(self.header['PT'])
        self.checksum = int(self.header['CS'])
        self.timestamp = int(mktime(strptime(self.header['FD'], '%m/%d/%Y %I:%M:%S %p')))
        self.payload = self.data[86:-6]
        if len(self.payload) != 1024:
            raise ValueError('Packet is the wrong size!')
        self.verify_checksum()

    def parse_header(self, data):
        if data[:6] != ('\x00' * 6):
            raise ValueError('Invalid packet header')
        data = data[6:]

        header = data.rstrip(' \r\n')
        header = header.split('/', 5)
        header = (x for x in header if x)
        header = ((x[:2], x[2:].strip(' ')) for x in header)
        return dict(header)

    def verify_checksum(self):
        checksum = sum([ord(x) for x in self.payload])
        if int(self.checksum) != checksum:
            raise ValueError('Checksum failed! Got: %i Expecting: %i' % (checksum, self.checksum))

    def dict(self):
        d = {}
        for field in ('filename', 'block', 'total_blocks', 'timestamp'):
            value = getattr(self, field)
            d[field] = value
        return d

    def __str__(self):
        return '%s (%i/%i)' % (self.filename, self.block, self.total_blocks)


class FileAssembler(object):
    def __init__(self, filename, callback=None):
        self.filename = filename
        self.callback = callback
        self.parts = {}

    def add_part(self, packet):
        self.parts[packet.block] = packet.payload
        self.check_parts(packet)

    def check_parts(self, packet):
        if self.callback is None:
            return

        if not None in [self.parts.get(i, None) for i in range(1, packet.total_blocks + 1)]:
            parts = self.parts.items()
            parts.sort(key=lambda x: x[0])
            content = ''.join([x[1] for x in parts])
            if not self.filename.endswith('.ZIS'):
                content = content.rstrip('\x00')
            self.content = content
            self.callback(self.filename, self.content)