def write(self, data):
if self.__written:
raise Exception("Internal Error. This should never happen")
self.__written = True
if not self.__deserialized:
self.__deserialized, desBytes = MessageData.Deserialize(data)
endTime = time.time()
runTimeInSeconds = int(endTime - self.__ctx.startTime)
runTimeInSeconds += 1
logger.info("Finished execution of code in %f seconds" %
runTimeInSeconds)
response = MessageData.GetMessageBuilder(EncryptedMobileCodeResult)
response["Cookie"].setData(self.__ctx.cookie)
response["RunTime"].setData(runTimeInSeconds)
response["RunMobileCodeHash"].setData(self.__ctx.runMobileCodeHash)
if self.__deserialized.topLevelData(
)[0] != definitions.playground.base.MobileCodeResult.PLAYGROUND_IDENTIFIER:
response["Success"].setData(False)
response["EncryptedResult"].setData("")
else:
encrypter = AES.new(self.__key, mode=AES.MODE_CBC, IV=self.__iv)
padder = playground.crypto.Pkcs7Padding(AES.block_size)
encrypted = encrypter.encrypt(padder.padData(data))
response["Success"].setData(self.__deserialized["success"].data())
response["EncryptedMobileCodeResultPacket"].setData(encrypted)
# in some ways, it would be easier to save "response" rather than
# response serialized. But we're saving this stuff to disk in case
# of interruption or disconnect. So serialized it is.
self.__ctx.finishCallback(self.__ctx, response.serialize())
packetTrace(logger, response,
"Encrypted mobile code result ready for transmission")
def fixedSizeTemplate(self, n, m):
bytes = ""
msgBuilder = MessageData.GetMessageBuilder(DummyMessage)
for i in range(1,n):
msgBuilder["data"].setData(str(i)*m)
bytes += Packet.MsgToPacketBytes(msgBuilder)
bufferOffset = 0
for i in range(1,n):
resultCode, result = Packet.RestorePacket(bytes, bufferOffset)
# check that we got a message
self.assertEqual(resultCode, Packet.BUFFER_STATUS_CONTAINS_MESSAGE,
"Message %d failed: %s" % (i, resultCode))
# deserialize the message
restoredData, bufferOffset = result
restoredMsgBuilder, bytesUsed = MessageData.Deserialize(restoredData)
# make sure the data matches
self.assertEqual(restoredMsgBuilder["data"].data(), str(i)*m,
"Restoration of message %d failed. Bad data" % i)
self.assertEqual(bytesUsed, len(restoredData),
"Didn't use up all bytes in message %d deserialization" % i)
# make sure that we used up all the bytes
self.assertEqual(bufferOffset, len(bytes),"Didn't use up all bytes in packet processing")
def encapsulatedMessageHandler(protocol, msg):
try:
c2cMsg, actualBytes = MessageData.Deserialize(msg.G2GMessage)
except:
# todo, this is an error. There wasn't a full message. Add logging?
return
if g_InterceptionDb.has_key(msg.Address) and g_InterceptionDb[
msg.Address].chaperoneProtocol == protocol:
key = c2cMsg.playground_msgID
messageInterceptor.ignore.add(key)
protocol.dataReceived(Packet.MsgToPacketBytes(c2cMsg))
def __handleRunMobileCode(self, prot, msg):
msgObj = msg.data()
self.__curState = self.__factory.getSessionState(msgObj.Cookie)
if not self.__curState or not self.__curState.state == self.STATE_OPEN:
curState = self.__curState and self.__curState.state or "<NO STATE>"
return self.__error(
"Invalid command. Cannot run mobile code unless session open (%s)"
% curState,
fatal=True)
logger.info("State found for cookie %s. State=%s" %
(msgObj.Cookie, self.__curState.state))
if msgObj.MaxRuntime > self.CODE_TIMEOUT:
response = MessageData.GetMessageBuilder(RunMobileCodeAck)
response["Cookie"].setData(self.__curState.cookie)
response["MobileCodeAccepted"].setData(False)
response["Message"].setData("Max Run Time parameter is too long.")
self.writeMsgAndClose(response)
rawRunMobileCodeMsg = msgObj.RunMobileCodePacket
startTime = time.time()
ctx = CodeExecutionContext()
ctx.startTime = startTime
ctx.cookie = self.__curState.cookie
ctx.runMobileCodeHash = SHA.new(rawRunMobileCodeMsg).digest()
ctx.finishCallback = lambda ctx, response: self.__factory.mobileCodeComplete(
ctx.cookie, response)
aesKey = os.urandom(16)
aesIv = os.urandom(16)
succeed, errmsg = self.__factory.createMobileCodeRecord(
self.__curState.cookie, aesKey, aesIv, msgObj.MaxRuntime)
if not succeed:
return self.__error("Could not run this code. Reason: " + errmsg,
fatal=True)
transport = WrapMobileCodeResultTransport(self.transport.getPeer(),
self.transport.getHost(),
aesKey, aesIv, ctx)
wrappedProtocol = WrapMobileCodeResultProtocol(transport)
logger.info("Starting execution of mobile code. MaxRunTime: %d" %
msgObj.MaxRuntime)
#realCodeHandler = playground.extras.sandbox.SandboxCodeunitAdapter(self.SANDBOX_CONTROLLER,
#timeout=min(msgObj.MaxRuntime,self.CODE_TIMEOUT))
#codeHandler = lambda codeUnit: self.__codeHandlerWrapper(realCodeHandler, codeUnit)
runMobileCodeHandler = RunMobileCodeHandler(
self, sandbox=SandboxCodeRunner())
runMobileCodeHandler(wrappedProtocol,
MessageData.Deserialize(rawRunMobileCodeMsg)[0])
self.__curState.state = self.STATE_RUNNING
response = MessageData.GetMessageBuilder(RunMobileCodeAck)
response["Cookie"].setData(self.__curState.cookie)
response["MobileCodeAccepted"].setData(True)
self.writeMsgAndClose(response)
def test_basicProtocol(self):
for messageSize in [0, 1, 100, 100000, 1000000]:
randomMessage = os.urandom(messageSize)
echoMessage = MessageData.GetMessageBuilder(EchoProtocolMessage)
echoMessage["original"].setData("True")
echoMessage["data"].setData(randomMessage)
self.clientTransport.writeMessage(echoMessage)
response = self.serverTransport.storage.pop()
responseMessageBuilder, bytesConsumed = MessageData.Deserialize(response)
responseData = responseMessageBuilder.data()
self.assertEqual(responseData.original, False, "Echo response should not be original")
self.assertEqual(responseData.data, randomMessage, "Echo message doesn't match")
self.assertEqual(len(response), bytesConsumed, "Message only took up %d bytes of %d byte response" % (bytesConsumed, len(response)))