pylinting round2

pull/10/head
Thorsten Sick 3 years ago
parent 033eb74886
commit ab6787241c

@ -6,12 +6,12 @@ import json
import os
import time
from pprint import pprint, pformat
import requests
import simplejson
from app.exceptions import CalderaError
from app.interface_sfx import CommandlineColors
from pprint import pprint, pformat
# TODO: Ability deserves an own class.
@ -51,7 +51,8 @@ class CalderaControl():
fullurl = self.url + "file/download"
request = requests.get(fullurl, headers=header)
filename = request.headers["FILENAME"] + extension
open(os.path.join(target_dir, filename), "wb").write(request.content)
with open(os.path.join(target_dir, filename), "wb") as fh:
fh.write(request.content)
# print(r.headers)
return filename
@ -221,6 +222,10 @@ class CalderaControl():
# ######### Get by id
def get_source(self, source_name):
""" Retrieves data source and detailed facts
@param: The name of the source
"""
payload = {"index": "sources",
"name": source_name}
@ -343,8 +348,8 @@ class CalderaControl():
facts = []
if parameters is not None:
for k, v in parameters.items():
facts.append({"trait": k, "value": v})
for key, value in parameters.items():
facts.append({"trait": key, "value": value})
payload["facts"] = facts
print(payload)
@ -368,8 +373,8 @@ class CalderaControl():
sources_name = "source_" + name
self.add_sources(sources_name, parameters)
print("Got:")
print(self.get_source("source_name"))
# To verify:
# print(self.get_source(sources_name))
payload = {"index": "operations",
"name": name,
@ -434,8 +439,8 @@ class CalderaControl():
facts = []
if parameters is not None:
for k, v in parameters.items():
facts.append({"trait": k, "value": v})
for key, value in parameters.items():
facts.append({"trait": key, "value": value})
payload["facts"] = facts
print(payload)

@ -208,8 +208,8 @@ class ExperimentConfig():
raise ConfigurationError("results missing in configuration")
try:
res = self.raw_config["results"]["loot_dir"]
except KeyError:
raise ConfigurationError("results/loot_dir not properly set in configuration")
except KeyError as error:
raise ConfigurationError("results/loot_dir not properly set in configuration") from error
return res
def attack_conf(self, attack):

@ -18,14 +18,16 @@ import os
class Metasploit():
def __init__(self, password, **kwargs):
def __init__(self, password, attack_logger, **kwargs):
"""
:param password: password for the msfrpcd
:param attack_logger: The attack logger to use for logging/printing
:param kwargs: Relevant ones: uri, port, server, username
"""
self.password = password
self.attack_logger = attack_logger
self.username = kwargs.get("username", None)
self.kwargs = kwargs
self.client = None
@ -353,8 +355,7 @@ class MetasploitInstant(Metasploit):
:param attack_logger: The attack logging
:param kwargs: Relevant ones: uri, port, server, username
"""
super().__init__(password, **kwargs)
self.attack_logger = attack_logger
super().__init__(password, attack_logger, **kwargs)
def parse_ps(self, ps_output):
d = []

@ -15,13 +15,13 @@ class SSHFeatures(BasePlugin):
def __init__(self):
super().__init__()
self.c = None
self.connection = None
def connect(self):
""" Connect to a machine """
if self.c:
return self.c
if self.connection is not None:
return self.connection
retries = 10
retry_sleep = 10
@ -31,7 +31,7 @@ class SSHFeatures(BasePlugin):
if self.config.os() == "linux":
uhp = self.get_ip()
self.vprint(f"Connecting to {uhp}", 3)
self.c = Connection(uhp, connect_timeout=timeout)
self.connection = Connection(uhp, connect_timeout=timeout)
if self.config.os() == "windows":
args = {}
@ -43,15 +43,15 @@ class SSHFeatures(BasePlugin):
self.vprint(args, 3)
uhp = self.get_ip()
self.vprint(uhp, 3)
self.c = Connection(uhp, connect_timeout=timeout, user=self.config.ssh_user(), connect_kwargs=args)
self.connection = Connection(uhp, connect_timeout=timeout, user=self.config.ssh_user(), connect_kwargs=args)
except (paramiko.ssh_exception.SSHException, socket.timeout):
self.vprint(f"Failed to connect, will retry {retries} times. Timeout: {timeout}", 0)
retries -= 1
timeout += 10
time.sleep(retry_sleep)
else:
self.vprint(f"Connection: {self.c}", 3)
return self.c
self.vprint(f"Connection: {self.connection}", 3)
return self.connection
self.vprint("SSH network error", 0)
raise NetworkError
@ -71,11 +71,12 @@ class SSHFeatures(BasePlugin):
self.vprint("Running SSH remote run: " + cmd, 3)
self.vprint("Disown: " + str(disown), 3)
# self.vprint("Connection: " + self.connection, 1)
result = None
retry = 2
while retry > 0:
try:
result = self.c.run(cmd, disown=disown)
result = self.connection.run(cmd, disown=disown)
print(result)
# paramiko.ssh_exception.SSHException in the next line is needed for windows openssh
except (paramiko.ssh_exception.NoValidConnectionsError, UnexpectedExit, paramiko.ssh_exception.SSHException):
@ -113,7 +114,7 @@ class SSHFeatures(BasePlugin):
timeout = 30
while retries:
try:
res = self.c.put(src, dst)
res = self.connection.put(src, dst)
except (paramiko.ssh_exception.SSHException, socket.timeout, UnexpectedExit):
self.vprint(f"PUT Failed to connect, will retry {retries} times. Timeout: {timeout}", 3)
retries -= 1
@ -144,7 +145,7 @@ class SSHFeatures(BasePlugin):
retry = 2
while retry > 0:
try:
res = self.c.get(src, dst)
res = self.connection.get(src, dst)
except (paramiko.ssh_exception.NoValidConnectionsError, UnexpectedExit):
if retry <= 0:
raise NetworkError
@ -163,6 +164,6 @@ class SSHFeatures(BasePlugin):
def disconnect(self):
""" Disconnect from a machine """
if self.c:
self.c.close()
self.c = None
if self.connection:
self.connection.close()
self.connection = None

@ -32,7 +32,7 @@ class FIN7Plugin(AttackPlugin):
if self.metasploit_1:
return self.metasploit_1
self.metasploit_1 = Metasploit(self.metasploit_password, attacker=self.attacker_machine_plugin, username=self.metasploit_user)
self.metasploit_1 = Metasploit(self.metasploit_password, attack_logger=self.attack_logger, attacker=self.attacker_machine_plugin, username=self.metasploit_user)
self.metasploit_1.start_exploit_stub_for_external_payload(payload=self.payload_type_1)
self.metasploit_1.wait_for_session()
return self.metasploit_1

@ -28,7 +28,7 @@ class VagrantPlugin(SSHFeatures, MachineryPlugin):
super().__init__()
self.plugin_path = __file__
self.v = None
self.c = None
self.connection = None
self.vagrantfilepath = None
self.vagrantfile = None
self.sysconf = {}
@ -80,13 +80,13 @@ class VagrantPlugin(SSHFeatures, MachineryPlugin):
# For linux we are using Vagrant style
if self.config.os() == "linux":
if self.c:
return self.c
if self.connection:
return self.connection
uhp = self.v.user_hostname_port(vm_name=self.config.vmname())
self.vprint(f"Connecting to {uhp}", 3)
self.c = Connection(uhp, connect_kwargs={"key_filename": self.v.keyfile(vm_name=self.config.vmname())})
return self.c
self.connection = Connection(uhp, connect_kwargs={"key_filename": self.v.keyfile(vm_name=self.config.vmname())})
return self.connection
else:
return super().connect()

@ -269,9 +269,6 @@ class TestExample(unittest.TestCase):
"relationships": [],
"facts": []
}
exp2 = {"index": "sources",
"name": "source_name"
}
exp3 = {"index": "operations",
"name": name,
"state": state,
@ -288,7 +285,7 @@ class TestExample(unittest.TestCase):
with patch.object(self.cc, "__contact_server__", return_value=None) as mock_method:
self.cc.add_operation(name, advid, group, state)
# mock_method.assert_called_once_with(exp, method="put")
mock_method.assert_has_calls([call(exp1, method="put"), call(exp2), call(exp3, method="put")])
mock_method.assert_has_calls([call(exp1, method="put"), call(exp3, method="put")])
# add_operation defaults
def test_add_operation_defaults(self):
@ -301,9 +298,6 @@ class TestExample(unittest.TestCase):
"relationships": [],
"facts": []
}
exp2 = {"index": "sources",
"name": "source_name"
}
exp3 = {"index": "operations",
"name": name,
"state": "running", # default
@ -319,7 +313,7 @@ class TestExample(unittest.TestCase):
}
with patch.object(self.cc, "__contact_server__", return_value=None) as mock_method:
self.cc.add_operation(name, advid)
mock_method.assert_has_calls([call(exp1, method="put"), call(exp2), call(exp3, method="put")])
mock_method.assert_has_calls([call(exp1, method="put"), call(exp3, method="put")])
# add_adversary
def test_add_adversary(self):

Loading…
Cancel
Save