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)