def main():
"""Connects to hidra and stores the streamed data to disk.
"""
parser = argparse.ArgumentParser()
parser.add_argument("--signal_host",
type=str,
help="Host where HiDRA is running",
default=socket.getfqdn())
parser.add_argument("--target_host",
type=str,
help="Host where the data should be send to",
default=socket.getfqdn())
arguments = parser.parse_args()
# enable logging
logfile_path = os.path.join(BASE_DIR, "logs")
logfile = os.path.join(logfile_path, "testAPI.log")
utils.init_logging(logfile, True, "DEBUG")
targets = [arguments.target_host, "50100", 0]
print("\n==== TEST: Stream all files and store them ====\n")
query = Transfer("STREAM", arguments.signal_host, use_log=True)
query.initiate(targets)
query.start()
target_dir = os.path.join(BASE_DIR, "data", "zmq_target")
target_file = os.path.join(target_dir, "test_store")
try:
query.store(target_file)
except Exception as excp:
print("Storing data failed.")
print("Error was:", excp)
query.stop()
print("\n==== TEST END: Stream all files and store them ====\n")
class DataReceiver(object):
"""Receives data and stores it to disc usign the hidra API.
"""
def __init__(self):
self.transfer = None
self.checking_thread = None
self.timeout = None
self.config = None
self.log = None
self.dirs_not_to_create = None
self.lock = None
self.target_dir = None
self.data_ip = None
self.data_port = None
self.transfer = None
self.checking_thread = None
self.plugin_handler = None
self.run_loop = True
self.setup()
self.exec_run()
def setup(self):
"""Initializes parameters, logging and transfer object.
"""
global _whitelist
try:
self.config = argument_parsing()
except Exception:
self.log = logging.getLogger("DataReceiver")
raise
config_gen = self.config["general"]
config_recv = self.config["datareceiver"]
# change user
user_info, user_was_changed = utils.change_user(config_gen)
# set up logging
utils.check_writable(config_gen["log_file"])
self._setup_logging()
utils.log_user_change(self.log, user_was_changed, user_info)
# set process name
# pylint: disable=no-member
setproctitle.setproctitle(config_gen["procname"])
self.log.info("Version: %s", __version__)
self.dirs_not_to_create = config_gen["dirs_not_to_create"]
# for proper clean up if kill is called
signal.signal(signal.SIGTERM, self.signal_term_handler)
self.timeout = 2000
self.lock = threading.Lock()
try:
ldap_retry_time = config_gen["ldap_retry_time"]
except KeyError:
ldap_retry_time = 10
try:
check_time = config_gen["netgroup_check_time"]
except KeyError:
check_time = 2
if config_gen["whitelist"] is not None:
self.log.debug("config_gen['whitelist']=%s",
config_gen["whitelist"])
with self.lock:
_whitelist = utils.extend_whitelist(config_gen["whitelist"],
config_gen["ldapuri"],
self.log)
self.log.info("Configured whitelist: %s", _whitelist)
else:
_whitelist = None
# only start the thread if a netgroup was configured
if (config_gen["whitelist"] is not None
and isinstance(config_gen["whitelist"], str)):
self.log.debug("Starting checking thread")
try:
self.checking_thread = CheckNetgroup(config_gen["whitelist"],
self.lock,
config_gen["ldapuri"],
ldap_retry_time,
check_time)
self.checking_thread.start()
except Exception:
self.log.error("Could not start checking thread",
exc_info=True)
else:
self.log.debug("Checking thread not started: %s",
config_gen["whitelist"])
self.target_dir = os.path.normpath(config_recv["target_dir"])
self.data_ip = config_recv["data_stream_ip"]
self.data_port = config_recv["data_stream_port"]
self.log.info("Writing to directory '%s'", self.target_dir)
self.transfer = Transfer(connection_type="STREAM",
use_log=True,
dirs_not_to_create=self.dirs_not_to_create)
self._load_plugin()
def _setup_logging(self):
config_gen = self.config["general"]
# enable logging
root = logging.getLogger()
root.setLevel(logging.DEBUG)
handlers = utils.get_log_handlers(config_gen["log_file"],
config_gen["log_size"],
config_gen["verbose"],
config_gen["onscreen"])
if isinstance(handlers, tuple):
for hdl in handlers:
root.addHandler(hdl)
else:
root.addHandler(handlers)
self.log = logging.getLogger("DataReceiver")
def _load_plugin(self):
try:
plugin_name = self.config["datareceiver"]["plugin"]
plugin_config = self.config[plugin_name]
except KeyError:
self.log.debug("No plugin specified")
return
self.plugin_handler = PluginHandler(plugin_name, plugin_config,
self.target_dir, self.log)
def exec_run(self):
"""Wrapper around run to react to exceptions.
"""
try:
self.run()
except KeyboardInterrupt:
pass
except Exception:
self.log.error("Stopping due to unknown error condition",
exc_info=True)
raise
finally:
self.stop()
def run(self):
"""Start the transfer and store the data.
"""
global _whitelist # pylint: disable=global-variable-not-assigned
global _changed_netgroup
if self.plugin_handler is not None:
plugin_type = self.plugin_handler.get_data_type()
self.plugin_handler.start()
else:
plugin_type = None
try:
self.transfer.start([self.data_ip, self.data_port], _whitelist)
except Exception:
self.log.error("Could not initiate stream", exc_info=True)
self.stop(store=False)
raise
# enable status check requests from any sender
self.transfer.setopt("status_check")
# enable confirmation reply if this is requested in a received data
# packet
self.transfer.setopt("confirmation")
self.log.debug("Waiting for new messages...")
self.run_loop = True
# run loop, and wait for incoming messages
while self.run_loop:
if _changed_netgroup:
self.log.debug("Re-registering whitelist")
self.transfer.register(_whitelist)
# reset flag
with self.lock:
_changed_netgroup = False
try:
ret_val = self.transfer.store(target_base_path=self.target_dir,
timeout=self.timeout,
return_type=plugin_type)
except KeyboardInterrupt:
break
except Exception:
self.log.error("Storing data...failed.", exc_info=True)
raise
if self.plugin_handler is None or ret_val is None:
continue
try:
self.plugin_handler.put(ret_val)
# ret_val might have been mutated by the plugin and therefore
# should only be reused if this is acceptable
except Exception:
self.log.error("Cannot submit message to plugin")
def stop(self, store=True):
"""Stop threads, close sockets and cleans up.
Args:
store (optional, bool): Run a little longer to store remaining
data.
"""
self.run_loop = False
if self.transfer is not None:
self.transfer.status = [b"ERROR", "receiver is shutting down"]
if store:
stop_timeout = 0.5
start_time = time.time()
diff_time = (time.time() - start_time) * 1000
self.log.debug("Storing remaining data.")
while diff_time < stop_timeout:
try:
self.log.debug("Storing remaining data...")
self.transfer.store(self.target_dir, self.timeout)
except Exception:
self.log.error("Storing data...failed.", exc_info=True)
diff_time = (time.time() - start_time) * 1000
self.log.info("Shutting down receiver...")
self.transfer.stop()
self.transfer = None
if self.plugin_handler is not None:
self.plugin_handler.stop()
if self.checking_thread is not None:
self.checking_thread.stop()
self.checking_thread.join()
self.log.debug("checking_thread stopped")
self.checking_thread = None
# pylint: disable=unused-argument
def signal_term_handler(self, signal_to_react, frame):
"""React on external SIGTERM signal.
"""
self.log.debug('got SIGTERM')
self.stop()
def __exit__(self, exception_type, exception_value, traceback):
self.stop()
def __del__(self):
self.stop()