def RunWithSendToSocket(self, dump_option, conditions=None):
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.bind((socket.gethostname(), 0))
port = sock.getsockname()[1]
send_to_socket_action = rdfvalue.MemoryCollectorSendToSocketAction(
host=socket.gethostname(),
port=port,
key=self.key,
iv=self.iv,
dump_option=dump_option)
flow_urn = flow.GRRFlow.StartFlow(
client_id=self.client_id,
flow_name="MemoryCollector",
conditions=conditions or [],
action=rdfvalue.MemoryCollectorAction(
action_type=rdfvalue.MemoryCollectorAction.Action.
SEND_TO_SOCKET,
send_to_socket=send_to_socket_action),
token=self.token,
output=self.output_path)
socket_data = []
def ReadFromSocket():
sock.listen(1)
client_socket, _ = sock.accept()
while 1:
data = client_socket.recv(1024)
if not data: break
socket_data.append(data)
client_socket.close()
sock.close()
thread = threading.Thread(target=ReadFromSocket)
thread.daemon = True
thread.start()
for _ in test_lib.TestFlowHelper(flow_urn,
self.client_mock,
client_id=self.client_id,
token=self.token):
pass
thread.join()
encrypted_data = "".join(socket_data)
# Data should be encrypted, so they're not equal
self.assertNotEqual(encrypted_data, self.memory_dump)
cipher = crypto.AES128CBCCipher(key=self.key,
iv=self.iv,
mode=crypto.AES128CBCCipher.OP_DECRYPT)
decrypted_data = cipher.Update(encrypted_data) + cipher.Final()
return flow_urn, encrypted_data, decrypted_data
def RunWithNoAction(self, conditions=None):
flow_urn = flow.GRRFlow.StartFlow(
client_id=self.client_id, flow_name="MemoryCollector",
conditions=conditions or [],
action=rdfvalue.MemoryCollectorAction(
action_type=rdfvalue.MemoryCollectorAction.Action.NONE),
token=self.token, output=self.output_path)
for _ in test_lib.TestFlowHelper(
flow_urn, self.client_mock, client_id=self.client_id, token=self.token):
pass
return aff4.FACTORY.Open(flow_urn, token=self.token)
def RunWithDownload(self, dump_option, conditions=None):
download_action = rdfvalue.MemoryCollectorDownloadAction(
dump_option=dump_option)
flow_urn = flow.GRRFlow.StartFlow(
client_id=self.client_id, flow_name="MemoryCollector",
conditions=conditions or [],
action=rdfvalue.MemoryCollectorAction(
action_type=rdfvalue.MemoryCollectorAction.Action.DOWNLOAD,
download=download_action
), token=self.token, output=self.output_path)
for _ in test_lib.TestFlowHelper(
flow_urn, self.client_mock,
client_id=self.client_id,
token=self.token):
pass
return aff4.FACTORY.Open(flow_urn, token=self.token)