Commit 9c27119a authored by Lukas Nemec's avatar Lukas Nemec

Fixed lint and tests.

parent 037113dd
# pylint: disable=missing-docstring
import traceback
import cachetools
import requests
import traceback
from haas_proxy import constants
class Balancer():
"""
Handles "load-balancing" of proxies between multiple running honeypots.
We call HTTP GET where we receive randomly assigned honeypot for 1H.
"""
api_url = None
# Expiring cache for API result, expires in 1h.
......@@ -16,10 +22,14 @@ class Balancer():
self.api_url = api_url
def load_api(self):
"""
Returns cached API response or get's data from API.
"""
cached_resp = self.cache.get(self.CACHE_KEY)
if cached_resp is None:
try:
resp = requests.api.get(self.api_url)
# pylint: disable=broad-except
except Exception:
traceback.print_exc()
return None
......@@ -34,6 +44,9 @@ class Balancer():
@property
def host(self):
"""
Returns host of honeypot.
"""
api_resp = self.load_api()
# load_api() may return None if there was error loading the API.
if api_resp is None:
......@@ -48,6 +61,9 @@ class Balancer():
@property
def port(self):
"""
Returns port of honeypot.
"""
api_resp = self.load_api()
# load_api() may return None if there was error loading the API.
if api_resp is None:
......
......@@ -109,14 +109,20 @@ class ProxySSHUser(ConchUser):
def getUserGroupId(self):
"""
Returns tuple with user and group ID.
Method needed by `SSHSessionForUnixConchUser`.
Method needed by `SSHSessionForUnixConchUser.openShell`.
"""
return 0, 0
def getHomeDir(self):
"""
Method needed by `SSHSessionForUnixConchUser.openShell`.
"""
return "/root"
def getShell(self):
"""
Method needed by `SSHSessionForUnixConchUser.openShell`.
"""
return "/bin/bash"
......
......@@ -24,16 +24,13 @@ def read_key(filename, default):
class Options(usage.Options):
optParameters = [
['device-token', 'd', None,
'Your ID at honeypot.labs.nic.cz. If you don\'t have one, sign up first.'],
['device-token', 'd', None, 'Your ID at honeypot.labs.nic.cz. If you don\'t have one, sign up first.'],
['port', 'p', constants.DEFAULT_PORT, 'Port to listen to.', int],
['balancer-address', None, constants.DEFAULT_BALANCER_ADDRESS],
['public-key'],
['private-key'],
['log-file', 'l', None,
'Turn on Python logging to this file. It\' wise to disable Twisted logging.'],
['log-level', None, 'warning',
'Possible options: error / warning / info / debug.'],
['log-file', 'l', None, 'Turn on Python logging to this file. It\' wise to disable Twisted logging.'],
['log-level', None, 'warning', 'Possible options: error / warning / info / debug.'],
]
@property
......
......@@ -30,7 +30,8 @@ def avatar(request):
username=request.param[0],
password=request.param[1],
)
avatar.conn.transport.transport.getPeer.return_value = mock.Mock(host='hacker', port=12345)
avatar.conn.transport.transport.getPeer.return_value = mock.Mock(
host='hacker', port=12345)
return avatar
......@@ -38,6 +39,7 @@ def avatar(request):
def proxy_ssh_session(cmd_args, avatar):
session = ProxySSHSession(avatar)
session.cmd_args = cmd_args
session.balancer = mock.Mock(host="localhost", port=2222)
return session
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment