From 4de90584c9db3f6154568a74f8c2451ac35e6816 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Mon, 28 Mar 2022 21:11:35 +0200 Subject: [PATCH] Agent: Add Credentials intercepting telemetry messenger --- ...ntials_intercepting_telemetry_messenger.py | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 monkey/infection_monkey/telemetry/messengers/credentials_intercepting_telemetry_messenger.py diff --git a/monkey/infection_monkey/telemetry/messengers/credentials_intercepting_telemetry_messenger.py b/monkey/infection_monkey/telemetry/messengers/credentials_intercepting_telemetry_messenger.py new file mode 100644 index 000000000..541800577 --- /dev/null +++ b/monkey/infection_monkey/telemetry/messengers/credentials_intercepting_telemetry_messenger.py @@ -0,0 +1,38 @@ +from functools import singledispatch + +from infection_monkey.credential_store import ICredentialsStore +from infection_monkey.telemetry.credentials_telem import CredentialsTelem +from infection_monkey.telemetry.i_telem import ITelem +from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger + + +class CredentialsInterceptingTelemetryMessenger(ITelemetryMessenger): + def __init__( + self, telemetry_messenger: ITelemetryMessenger, credentials_store: ICredentialsStore + ): + self._telemetry_messenger = telemetry_messenger + self._credentials_store = credentials_store + + def send_telemetry(self, telemetry: ITelem): + _send_telemetry(telemetry, self._telemetry_messenger, self._credentials_store) + + +# Note: We can use @singledispatchmethod instead of @singledispatch if we migrate to Python 3.8 or +# later. +@singledispatch +def _send_telemetry( + telemetry: ITelem, + telemetry_messenger: ITelemetryMessenger, + credentials_store: ICredentialsStore, +): + telemetry_messenger.send_telemetry(telemetry) + + +@_send_telemetry.register +def _( + telemetry: CredentialsTelem, + telemetry_messenger: ITelemetryMessenger, + credentials_store: ICredentialsStore, +): + credentials_store.add_credentials(telemetry.credentials) + telemetry_messenger.send_telemetry(telemetry)