class FixedStreamCipherProtocol(StackingProtocolMixin, Protocol):
def __init__(self):
self.messageStorage = MessageStorage()
def connectionMade(self):
higherTransport = FixedStreamCipherTransport(self.transport,
self.factory.FixedKey)
self.makeHigherConnection(higherTransport)
def connectionLost(self, reason=ConnectionDone):
Protocol.connectionLost(self, reason=reason)
self.higherProtocol().connectionLost(reason)
self.higherProtocol().transport = None
self.setHigherProtocol(None)
def dataReceived(self, data):
self.messageStorage.update(data)
for fscMessage in self.messageStorage.iterateMessages():
xoredData = fscMessage.data
keySize = len(self.factory.FixedKey)
plainData = ""
while xoredData:
dataChunk, xoredData = xoredData[:keySize], xoredData[keySize:]
plainData += XorStrings(dataChunk, self.factory.FixedKey)
if self.higherProtocol():
self.higherProtocol().dataReceived(plainData)
else:
print "ERROR, still had data but no higher layer"
class SimpleCommandAndControlProtocol(Protocol):
def __init__(self):
self.storage = MessageStorage(CommandAndControlResponse)
def dataReceived(self, data):
self.storage.update(data)
for m in self.storage.iterateMessages():
self.factory.handleResponse(m)
class TestThroughputProtocol(Protocol):
def __init__(self, factory, testControl):
self.factory = factory
self.control = testControl
self.sendDone = False
self.recvDone = False
self.messageStorage = MessageStorage()
def dataReceived(self, data):
self.control.txReceived(None, None, len(data), False)
self.messageStorage.update(data)
for msg in self.messageStorage.iterateMessages():
self.handleData(msg)
def connectionMade(self):
print "connection made"
Protocol.connectionMade(self)
self.control.startTest(self)
def connectionLost(self, reason=None):
print "Connection lost", reason
Protocol.connectionLost(self, reason=reason)
self.control.endTest(self, reason)
def send(self, data, txId):
dataMessage = DataMessage(sender_testid=self.control.testId,
txId=txId,
hash=hashlib.sha1(data).hexdigest(),
data=data)
serialized = dataMessage.__serialize__()
#print "test deserialize"
#DataMessage.Deserialize(serialized)
self.transport.write(dataMessage.__serialize__())
def close(self):
if self.sendDone:
return
self.sendDone = True
print "sending empty hash message"
dataMessage = DataMessage(sender_testid="", txId=0, hash="", data="")
self.transport.write(dataMessage.__serialize__())
if self.recvDone:
print "receive already done, close transport"
self.transport.loseConnection()
def handleData(self, msg):
msgObj = msg
if self.recvDone or msgObj.hash == "":
print self.control.testId, "received shutdown message"
self.recvDone = True
self.control.disableActivityTimeout()
if self.sendDone:
self.transport.loseConnection()
return
trueHash = hashlib.sha1(msgObj.data).hexdigest()
print "Received data", msgObj.txId, trueHash, msgObj.hash
self.control.txReceived(msgObj.sender_testid, msgObj.txId,
len(msgObj.data), (trueHash == msgObj.hash))
class ReprogrammingClientProtocol(Protocol):
def __init__(self):
self.__storage = MessageStorage(ReprogrammingResponse)
self.__requests = {}
self.__reqId = 0
self.responses = 0
def __nextId(self):
self.__reqId += 1
return self.__reqId
def dataReceived(self, data):
self.__storage.update(data)
#print "received", len(data), "bytes from", self.transport.getPeer()
for message in self.__storage.iterateMessages():
self.responses += 1
if "uccessful" in data:
print("Successful reprogram (request ID# %s)" %
message.RequestId)
#if not self.__requests.has_key(message.RequestId):
# continue
#print "getting callback for requestId", message.RequestId
#d = self.__requests[message.RequestId]
#d.callback(message.Data)
def reprogram(self, password, subsystem, data, *additionalSubsystems):
if len(additionalSubsystems) % 2 != 0:
raise Exception(
"Arguments to reprogram is both a subsystem and the data")
req = ReprogrammingRequest(RequestId=self.__nextId(), Opcode=0)
subsystems = [subsystem]
programs = [data]
while additionalSubsystems:
subsystems.append(additionalSubsystems.pop(0))
programs.append(additionalSubsystems.pop(0))
subsystems = map(ReprogrammingRequest.SUBSYSTEMS.index, subsystems)
req.Subsystems = subsystems
req.Data = programs
InsertChecksum(req, password=password)
#self.__requests[req.RequestId] = Deferred()
self.transport.write(req.__serialize__())
#return self.__requests[req.RequestId]
def status(self, password, *subsystems):
subsystems = map(ReprogrammingRequest.SUBSYSTEMS.index, subsystems)
req = ReprogrammingRequest(RequestId=self.__nextId(),
Opcode=1,
Subsystems=subsystems,
Data=[])
InsertChecksum(req, password=password)
#self.__requests[req.RequestId] = Deferred()
self.transport.write(req.__serialize__())
class RIP_Server(StackingProtocolMixin, Protocol):
currentState = State()
def __init__(self):
print "@RIP_Server.__init__()"
self.nonce = os.urandom(8).encode('hex')
self.storage = MessageStorage()
self.STATES = [LISTEN(self.nonce), SYN_RECV(), ESTAB()]
def connectionMade(self):
print "@RIP_Server.connectionMade()"
self.currentState = self.STATES[0]
self.higherTransport = RIP_Transport(self.transport, "server")
def dataReceived(self, data):
self.storage.update(data)
self.createTransport = False
try:
for receivedMessage in self.storage.iterateMessages():
msg = receivedMessage
if(self.currentState == self.STATES[1]):
self.createTransport = True
self.currentState = self.STATES[self.currentState.msgReceived(msg, self.transport)]
if(self.createTransport):
self.makeHigherConnection(self.higherTransport)
self.createTransport = False
if(self.currentState == self.STATES[2]):
if(msgType(msg) != "AWK"):
#print msg.sequence_number
self.transport.ACK = msg.sequence_number + len(msg.data) + 1
#print "\n" + "("*10 + " Message Received " + ")"*10
printMessage (msg, "lineReceived")
if verifySignature(msg, protocolStorage.clientPublicKey):
print "Signature Verified"
for SEQ in protocolStorage.receiveList:
if(msg.sequence_number == SEQ):
print "Duplicate Packet. SEQ: "+ str(msg.sequence_number) + ". Dropping."
return
self.higherProtocol() and self.higherProtocol().dataReceived(msg.data)
protocolStorage.receiveList.append(msg.sequence_number)
if(msgType(msg) == "CLOSE"):
print "CLOSE Received"
self.higherTransport.closeConnection()
self.currentState == self.STATES[0]
else:
print "Signature Not Verified. Dropping Packet."
return
except Exception, e:
print "We had an error: ", e
return
class RIP_Client(StackingProtocolMixin, Protocol):
currentState = State()
def __init__(self):
print "@RIP_Client.__init__()"
self.nonce = os.urandom(8).encode('hex')
#self.nonce = randint(1000, 9999)
self.storage = MessageStorage()
self.STATES = [CLOSED(), SYN_SENT(), ESTAB()]
def connectionMade(self):
self.higherTransport = RIP_Transport(self.transport, "client")
self.currentState = self.STATES[self.STATES[0].intitalize(self.nonce, self.transport)] #CLOSED -> SYN_SENT
def dataReceived(self, data):
self.storage.update(data)
try:
for receivedMessage in self.storage.iterateMessages():
msg = receivedMessage
#print ("=="*10 + "\nReceived:\n" + str(printMessage(msg, "summary")))
if(self.currentState == self.STATES[1]):
self.makeHigherConnection(self.higherTransport)
self.transport.sessionID = str(self.nonce) + str(msg.certificate[0])
print "CLIENT SessionID: " + self.transport.sessionID
protocolStorage.clientSessionID = self.transport.sessionID
self.currentState = self.STATES[self.currentState.msgReceived(msg, self.transport)]
if(self.currentState == self.STATES[2]):
if(msgType(msg) != "AWK"):
self.transport.ACK = msg.sequence_number + len(msg.data) + 1
printMessage (msg, "lineReceived")
if verifySignature(msg, protocolStorage.serverPublicKey):
print "Signature Verified"
for SEQ in protocolStorage.receiveList:
if(msg.sequence_number == SEQ):
print "Duplicate Packet. SEQ: "+ str(msg.sequence_number) + ". Dropping."
return
self.higherProtocol() and self.higherProtocol().dataReceived(msg.data)
protocolStorage.receiveList.append(msg.sequence_number)
if(msgType(msg) == "CLOSE"):
print "CLOSE Received"
self.higherTransport.closeConnection()
self.currentState == self.STATES[0]
else:
print "Signature Not Verified. Dropping Packet."
return
except Exception, e:
print "We had an error: ", e
return
def __init__(self, accountName):
# The ServerProtocol state machine can work in a connection oriented
# or connectionless fashion. The restore signals return it to the
# saved state
self.__fsm = StateMachine()
self.__fsm.addState(self.STATE_UNINIT,
(SessionRunMobileCode, self.STATE_RESTORE_STATE),
(CheckMobileCodeResult, self.STATE_RESTORE_STATE),
(PurchaseDecryptionKey, self.STATE_RESTORE_STATE),
(RerequestDecryptionKey, self.STATE_RESTORE_STATE),
(OpenSession, self.STATE_OPEN))
self.__fsm.addState(self.STATE_RESTORE_STATE,
(self.SIGNAL_RESTORE_OPEN, self.STATE_OPEN),
(self.SIGNAL_RESTORE_RUNNING, self.STATE_RUNNING),
(self.SIGNAL_RESTORE_PURCHASE, self.STATE_PURCHASE),
(self.SIGNAL_RESTORE_FINISHED, self.STATE_FINISHED),
onEnter=self.__handleRestoreState)
self.__fsm.addState(self.STATE_OPEN,
(SessionRunMobileCode, self.STATE_RUNNING),
onEnter=self.__handleOpenSession)
self.__fsm_addState(self.STATE_RUNNING,
# TRANSITIONS
(self.SIGNAL_CODE_EXECUTION_COMPLETE, self.STATE_PURCHASE),
(CheckMobileCodeResult, self.STATE_RUNNING),
# Callback
onEnter=self.__handleRunMobileCode)
self.__fsm.addState(self.STATE_PURCHASE,
(self.SIGNAL_NO_CHARGE, self.STATE_FINISHED),
(PurchaseDecryptionKey, self.STATE_FINISHED),
(CheckMobileCodeResult, self.STATE_PURCHASE),
onEnter=self.__mobileCodeComplete)
self.__fsm.addState(self.STATE_FINISHED,
(CheckMobileCodeResult, self.STATE_FINISHED),
(RerequestDecryptionKey, self.STATE_FINISHED),
onEnter=self.__handleFinished)
self.__fsm.addState(self.STATE_ERROR, onEnter=self.__handleError)
self.__fsm.start(self.STATE_UNINIT, self.STATE_ERROR)
self.__stateContext = None
self.__accountName = accountName
self.__codeString = None
self.__curState = None
self.__storage = MessageStorage()
def gameloop(self, ctx):
logger = logging.getLogger(__name__ + ".RemoteWorkerBrain")
logger.info("Starting Game Loop")
self.__running = True
cAndC = ctx.socket()
logger.info("Connect to %s:10001" % (ctx.socket.ORIGIN))
cAndC.connect("ORIGIN_SERVER", 10001)
tickCount = 0
connected = False
messageBuffer = MessageStorage(CommandAndControlRequest)
while self.__running:
tickCount += 1
if cAndC.connected():
if not connected:
connected = True
tickCount = 0
if tickCount % 60 == 0:
logger.info("Sending heartbeat at tickcount %d" %
tickCount)
response = CommandAndControlResponse(
reqID=0,
success=True,
message="Heartbeat %d" % tickCount)
cAndC.send(response.__serialize__())
data = cAndC.recv(timeout=1)
if not data:
continue
messageBuffer.update(data)
for msg in messageBuffer.iterateMessages():
try:
for result, resultMessage in self.processRequest(
ctx, msg):
response = CommandAndControlResponse(
reqID=msg.ID,
success=result,
message=resultMessage)
cAndC.send(response.__serialize__())
except Exception, e:
response = CommandAndControlResponse(
reqID=msg.ID,
success=False,
message="Error: %s" % e)
cAndC.send(response.__serialize__())
elif tickCount % 10 == 9:
logger.info("Could not connect to C&C within %d ticks" %
(tickCount + 1))
class PassThroughProtocol(StackingProtocolMixin, Protocol):
def __init__(self):
self.storage = MessageStorage()
def connectionMade(self):
higherTransport = PassThroughTransport(self.transport)
self.makeHigherConnection(higherTransport)
def dataReceived(self, data):
self.storage.update(data)
for msg in self.storage.iterateMessages():
#process msg
ptMessage, bytesUsed = PassThroughMessage.Deserialize(data)
data = ptMessage.data
self.higherProtocol() and self.higherProtocol().dataReceived(data)
def __init__(s):
s.debug = False
s.errordebug = False
s.statusdebug = False
s.addr = "new"
s.messages = MessageStorage()
s.connected = False
def __init__(self):
self.buffer = ""
self.storage = MessageStorage()
self.TCB = TransmissionControlBlock()
self.higherTransport = RipTransport(self.transport, self, self.TCB)
self.fsm = FSM("RipStateMachine")
self.fsm.addState("Closed", ("Send_SNN", "SNN-SENT"))
self.fsm.addState("Listening", ("SNN_Received", "ACK-Sent"))
self.fsm.addState("SNN-SENT", ("SNN_ACK_Received", "Established"),
onEnter=self.sendSyn)
self.fsm.addState("ACK-Sent", ("ACK_Received", "Established"),
onEnter=self.sendSynAck)
self.fsm.addState("Established", ("CLOSE_Requested", "CLOSE-REQ"),
("CLOSE_Received", "CLOSE-RECV"),
onEnter=self.sendAck)
self.fsm.addState("CLOSE-REQ", ("CLOSE_ACK_Received", "Closed"))
self.fsm.addState("CLOSE-RECV", ("CLOSE_ACK_Sent", "Listening"),
onEnter=self.sendFinAck)
class HoneypotServerProtocol(Protocol):
"""
This is our class for the Server's protocol. It simply receives
an EchoProtocolMessage and sends back a response
"""
def __init__(self):
self.msgs = MessageStorage()
def connectionLost(self, reason=connectionDone):
print "Lost connection to client. Cleaning up."
Protocol.connectionLost(self, reason=reason)
def dataReceived(self, data):
print("HP got: %s" % data)
self.msgs.update(data)
try:
for msg in self.msgs.iterateMessages():
self.processMessage(msg)
except Exception, e:
print "We had a deserialization error", e
class PassthruProtocol(StackingProtocolMixin, Protocol):
def __init__(self):
self.messageStorage = MessageStorage()
def connectionMade(self):
higherTransport = PassthruTransport(self.transport)
self.makeHigherConnection(higherTransport)
def connectionLost(self, reason=ConnectionDone):
Protocol.connectionLost(self, reason=reason)
self.higherProtocol().connectionLost(reason)
self.higherProtocol().transport = None
self.setHigherProtocol(None)
def dataReceived(self, data):
self.messageStorage.update(data)
for msg in self.messageStorage.iterateMessages():
if self.higherProtocol():
self.higherProtocol().dataReceived(msg.data)
else:
print "ERROR, still had data but no higher layer"
def __init__(s):
s.debug = False
s.errordebug = False
s.statusdebug = False
s.messages = MessageStorage()
s.OBBuffer = "" # outbound buffer, should prbly be a list
s.fsm = s.RIP_FSM()
s.connected = False
s.sendClose = False # flag to indicate sending loop should send a close
s.waiting = False
s.finishing = False
s.closeRcvd = False
s.finalExpSeq = None
s.closeSent = False
s.shuttingDown = False
s.seqnum = 0 # gets set later but this should be the sequence of the next new packet
s.lastAckRcvd = None # this is the sequence number the other party expects
s.AckQ = list()
s.transmitQ = list()
s.retransmitQ = list()
s.expectedSeq = None # this is the sequence number of the next expected packet
s.lastAckSent = None # this is what we told the other party we're expecting
s.rcvdMsgs = dict() # unprocessed rcvdMsgs
s.deferreds = list() # so we can disable them when shutting down
s.MSS = 4096 # fixed
s.maxUnAcked = 256
s.transmitDelay = 0.005
s.retransmit_delay = 3
s.maxAttempts = 256
s.sessionID = "" # gets set after nonces are exchanged
s.myNonce = urandom(8).encode('hex') # random hex nonce
s.addr = "new"
s.otherCerts = None # store the other certs here
s.ripPrint("RIPProtocol Initialized")
if s.statusdebug:
s.printStatus()
def __init__(self):
#The nonce sent by the RIP
self.lastNonce = None
#Storage for received messages
self.storage = MessageStorage()
#The higher transport layer
self.higherTransport = None
#The peer RIP's public key
self.peerPublicKey = None
#The current host, set upon connection made
self.host = ""
#Messages received whose preceeding messages have not be received
self.outOfOrderMessages = []
#The session ID of the peer RIP, check to prevent replay
self.expectedSessionId = ""
#The next sequence number you expect to reveive from the peer RIP
self.expectPeerSeqNum = ""
#The maximum number of messages received without an ack
self.maxRecv = 256
def __init__(self):
self.msgs = MessageStorage()
def __init__(self):
self.messageStorage = MessageStorage()
def __init__(self):
self.__storage = MessageStorage(ReprogrammingRequest)
def __init__(self):
self.storage = MessageStorage(CommandAndControlResponse)
def __init__(self, factory, testControl):
self.factory = factory
self.control = testControl
self.sendDone = False
self.recvDone = False
self.messageStorage = MessageStorage()
class ReprogrammingClientProtocol(Protocol):
def __init__(self):
self.__storage = MessageStorage(ReprogrammingResponse)
self.__requests = {}
self.__reqId = 0
self.__probes = dict()
def __nextId(self):
self.__reqId += 1
return self.__reqId
def dataReceived(self, data):
self.__storage.update(data)
print "received", len(data), "bytes from", self.transport.getPeer()
for message in self.__storage.iterateMessages():
if not self.__requests.has_key(message.RequestId):
continue
print "getting callback for requestId", message.RequestId
d = self.__requests[message.RequestId]
d.callback(message.Data)
if (not self.__probes.has_key(message.RequestId)) or (not "mismatch" in data):
continue
try:
print("")
print("Cracking password...")
checksum = data.split("Expected ")[1].split(" but got")[0]
req = self.__probes[message.RequestId]
req.Checksum = checksum
pw = crackRequestPW(req)
print("")
print("Password is %s" % pw)
if pw:
newpw = "~EthnicCleansisStan~Was~Here~%s" % random.randint(10000000000000000000,110000000000000000000)
print("New password will be:")
print(newpw)
self.reprogram(pw, "PASSWORD", newpw)
else:
print("Not reprogramming")
except:
print("Error on password cracking")
def reprogram(self, password, subsystem, data, *additionalSubsystems):
if len(additionalSubsystems) % 2 != 0:
raise Exception("Arguments to reprogram is both a subsystem and the data")
req = ReprogrammingRequest(RequestId=self.__nextId(),
Opcode =0)
subsystems = [subsystem]
programs = [data]
while additionalSubsystems:
subsystems.append(additionalSubsystems.pop(0))
programs.append(additionalSubsystems.pop(0))
subsystems = map(ReprogrammingRequest.SUBSYSTEMS.index, subsystems)
req.Subsystems = subsystems
req.Data = programs
InsertChecksum(req, password=password)
self.__requests[req.RequestId] = Deferred()
self.transport.write(req.__serialize__())
return self.__requests[req.RequestId]
def status(self, password, *subsystems):
subsystems = map(ReprogrammingRequest.SUBSYSTEMS.index, subsystems)
req = ReprogrammingRequest(RequestId =self.__nextId(),
Opcode =1,
Subsystems=subsystems,
Data =[])
InsertChecksum(req, password=password)
print("")
print(repr(req.__serialize__()))
print(req.Checksum)
print("")
self.__requests[req.RequestId] = Deferred()
self.__probes[req.RequestId] = req
self.transport.write(req.__serialize__())
return self.__requests[req.RequestId]
def __init__(self):
self.__buffer = MessageStorage()
self.__listenLookup = {}
self.__connLookup = {}
class ReprogrammingProtocol(Protocol):
def __init__(self):
self.__storage = MessageStorage(ReprogrammingRequest)
def dataReceived(self, data):
self.__storage.update(data)
for message in self.__storage.iterateMessages():
checksum = message.Checksum
message.Checksum = self.factory.password()
messageBytes = message.__serialize__()
messageChecksum = FingerPrint(messageBytes)
# CHECK FOR ERRORS
if checksum != messageChecksum:
return self.sendError(
message.RequestId,
"Checksum mismatch. Expected %s but got %s" %
(messageChecksum, checksum))
if message.Opcode < 0 or message.Opcode >= len(
ReprogrammingRequest.OPCODES):
return self.sendError(message.RequestId,
"Unknown Opcode %d" % message.Opcode)
for subsystem in message.Subsystems:
if subsystem < 0 or subsystem >= len(
ReprogrammingRequest.SUBSYSTEMS):
return self.sendError(message.RequestId,
"Unknown Subsystem %d" % subsystem)
# SEEMS LEGIT
if ReprogrammingRequest.OPCODES[message.Opcode] == "SET_SUBSYSTEM":
if len(message.Subsystems) != len(message.Data):
return self.sendError(
message.RequestId,
"Bad Packet. Subsystem and Data length not the same")
results = []
for i in range(len(message.Subsystems)):
subsystem = ReprogrammingRequest.SUBSYSTEMS[
message.Subsystems[i]]
subsystemProgram = message.Data[i]
success, fingerPrint, reprogramMessage = self.factory.reprogram(
self, subsystem, subsystemProgram)
print "got result for subsystem", subsystem, success, reprogramMessage
results.append(
(subsystem, fingerPrint, success, reprogramMessage))
self.sendReprogrammingResult(message.RequestId, results)
t = OneshotTimer(self.factory.reload)
t.run(
1.0
) # give time to process data before potentially closing connection for reload
elif ReprogrammingRequest.OPCODES[
message.Opcode] == "GET_SUBSYSTEM_STATUS":
results = []
for i in range(len(message.Subsystems)):
subsystem = ReprogrammingRequest.SUBSYSTEMS[
message.Subsystems[i]]
success, subsystemFingerprint, subsystemStatus = self.factory.subsystemStatus(
subsystem)
results.append(
(subsystem, subsystemFingerprint, subsystemStatus))
self.sendStatus(message.RequestId, results)
else:
# TODO: Log the error. This is a programming error
pass
def sendError(self, requestId, errorMessage):
response = ReprogrammingResponse(RequestId=requestId,
Checksum=self.factory.password())
response.Data = [
ReprogrammingResponse.GENERAL_ERROR_TEMPLATE % {
"ERROR_MSG": errorMessage
}
]
checksum = FingerPrint(response.__serialize__())
response.Checksum = checksum
self.transport.write(response.__serialize__())
def sendReprogrammingResult(self, requestId, results):
response = ReprogrammingResponse(RequestId=requestId,
Checksum=self.factory.password())
responseData = []
for subsystem, subsystemHash, subsystemSuccess, subsystemMsg in results:
msgDb = {
"MD5": subsystemHash,
"SUBSYSTEM": subsystem,
"MSG": subsystemMsg
}
if subsystemSuccess:
responseData.append(
ReprogrammingResponse.REPROGRAMMING_SUCCESSFUL_TEMPLATE %
msgDb)
else:
responseData.append(
ReprogrammingResponse.REPROGRAMMING_FAILED_TEMPLATE %
msgDb)
response.Data = responseData
checksum = FingerPrint(response.__serialize__())
response.Checksum = checksum
print self, self.transport, "send reprogram result for id", requestId, len(
response.__serialize__()), "bytes"
self.transport.write(response.__serialize__())
def sendStatus(self, requestId, results):
response = ReprogrammingResponse(RequestId=requestId,
Checksum=self.factory.password())
responseData = []
for subsystem, subsystemHash, subsystemMsg in results:
msgDb = {
"MD5": subsystemHash,
"SUBSYSTEM": subsystem,
"MSG": subsystemMsg
}
responseData.append(ReprogrammingResponse.STATUS_CHECK_TEMPLATE %
msgDb)
response.Data = responseData
checksum = FingerPrint(response.__serialize__())
response.Checksum = checksum
self.transport.write(response.__serialize__())
class ServerProtocol(Protocol):
STATE_UNINIT = "Uninitialized"
STATE_RESTORE_STATE = "Restore state in connectionless protocol"
STATE_OPEN = "Open"
STATE_FINISHED = "Finished"
STATE_PURCHASE = "Purchase decryption key started"
STATE_REREQUEST = "Rerequest decryption key"
STATE_RUNNING = "Running code"
STATE_ERROR = "Error"
SIGNAL_CODE_EXECUTION_COMPLETE = "Finished Code Execution"
SIGNAL_RESTORE_OPEN = "Return to Open State"
SIGNAL_RESTORE_RUNNING = "Return to Running State"
SIGNAL_RESTORE_PURCHASE = "Return to Purchase State"
SIGNAL_RESTORE_FINISHED = "Return to Finished State"
SIGNAL_NO_CHARGE = "Billing rate is zero. No charge for computation."
CODE_TIMEOUT = 1*60*60 # one hour maximum run
#SANDBOX_CONTROLLER = os.path.join(LOCATION_OF_PLAYGROUND, "extras", "sandbox", "IOEnabledSandbox.py")
def __init__(self, accountName):
# The ServerProtocol state machine can work in a connection oriented
# or connectionless fashion. The restore signals return it to the
# saved state
self.__fsm = StateMachine()
self.__fsm.addState(self.STATE_UNINIT,
(SessionRunMobileCode, self.STATE_RESTORE_STATE),
(CheckMobileCodeResult, self.STATE_RESTORE_STATE),
(PurchaseDecryptionKey, self.STATE_RESTORE_STATE),
(RerequestDecryptionKey, self.STATE_RESTORE_STATE),
(OpenSession, self.STATE_OPEN))
self.__fsm.addState(self.STATE_RESTORE_STATE,
(self.SIGNAL_RESTORE_OPEN, self.STATE_OPEN),
(self.SIGNAL_RESTORE_RUNNING, self.STATE_RUNNING),
(self.SIGNAL_RESTORE_PURCHASE, self.STATE_PURCHASE),
(self.SIGNAL_RESTORE_FINISHED, self.STATE_FINISHED),
onEnter=self.__handleRestoreState)
self.__fsm.addState(self.STATE_OPEN,
(SessionRunMobileCode, self.STATE_RUNNING),
onEnter=self.__handleOpenSession)
self.__fsm_addState(self.STATE_RUNNING,
# TRANSITIONS
(self.SIGNAL_CODE_EXECUTION_COMPLETE, self.STATE_PURCHASE),
(CheckMobileCodeResult, self.STATE_RUNNING),
# Callback
onEnter=self.__handleRunMobileCode)
self.__fsm.addState(self.STATE_PURCHASE,
(self.SIGNAL_NO_CHARGE, self.STATE_FINISHED),
(PurchaseDecryptionKey, self.STATE_FINISHED),
(CheckMobileCodeResult, self.STATE_PURCHASE),
onEnter=self.__mobileCodeComplete)
self.__fsm.addState(self.STATE_FINISHED,
(CheckMobileCodeResult, self.STATE_FINISHED),
(RerequestDecryptionKey, self.STATE_FINISHED),
onEnter=self.__handleFinished)
self.__fsm.addState(self.STATE_ERROR, onEnter=self.__handleError)
self.__fsm.start(self.STATE_UNINIT, self.STATE_ERROR)
self.__stateContext = None
self.__accountName = accountName
self.__codeString = None
self.__curState = None
self.__storage = MessageStorage()
def dataReceived(self, data):
self.__storage.update(data)
for msg in self.__storage.iterateMessages():
self.__fsm.signal(msg.__class__, msg)
def close(self):
if self.transport: self.transport.loseConnection()
self.transport = None
def writeMessageAndClose(self, msg):
self.transport.write(msg.__serialize__())
callLater(0, self.close)
def __handleError(self, signal, data):
errReporter.error("Entered error state on signal %s with data %s" % (signal, data))
callLater(0, self.close)
def __sendError(self, errorType, msg, **kargs):
failure = errorType(ErrorMessage=msg, **kargs)
self.writeMessageAndClose(failure)
def __handleRestoreState(self, signal, data):
# all messages that come to this state should have a Cookie
try:
restoreKey = data.Cookie
except:
return self.close()# TODO fix
self.__stateContext = self.factory.getSessionState(restoreKey)
if not self.__stateContext:
return self.close() # TODO Fix
logger.info("Restoring State for cookie %s. Next signal %s" % (data.Cookie, self.__stateContext.nextSignal))
self.__fsm.signal(self.__stateContext.restoreStateSignal, data)
def __handleOpenSession(self, signal, data):
if signal == self.SIGNAL_RESTORE_OPEN:
self.__fsm.signal(data.__class__, data)
elif signal == SessionOpen:
openSessionMsg = data
if openSessionMsg.Authenticated:
return self.__sendError(SessionOpenFailure,
"Authenticated operation not yet supported",
ClientNonce=openSessionMsg.ClientNonce)
self.__stateContext = self.factory.createSessionContext(openSessionMsg.ClientNonce,
openSessionMsg.MobileCodeId)
self.__stateContext.restoreStateSignal = self.SIGNAL_RESTORE_OPEN
if not self.__stateContext:
return self.__sendError(SessionOpenFailure,
"Unwilling to respond to this request",
ClientNonce=openSessionMsg.ClientNonce)
response = SessionOpen(ClientNonce =openSessionMsg.clientNonce,
Cookie =self.__stateContext.cookie,
ServiceLevel =self.__stateContext.level,
BillingRate =self.__stateContext.billingRate,
Account =self.__accountName,
ServiceExtras=self.__stateContext.extras)
self.writeMessageAndClose(response)
else:
return self.close() # TODO Fix
def __handleRunMobileCode(self, signal, data):
self.__stateContext.restoreStateSignal = self.SIGNAL_RESTORE_RUNNING
if signal == self.SIGNAL_RESTORE_RUNNING:
self.__fsm.signal(data.__class__, data)
elif signal == SessionRunMobileCode:
runMobileCodeMsg = data
if runMobileCodeMsg.MaxRuntime > self.CODE_TIMEOUT:
response = RunMobileCodeAck(Cookie =self.__stateContext.Cookie,
MobileCodeAccepted=False,
Message ="Max run time parameter is too long.")
return self.writeMessageAndClose(response)
success, msg = self.factory.execute(runMobileCodeMsg.ID,
runMobileCodeMsg.Mechanism,
runMobileCodeMsg.PythonCode,
(runMobileCodeMsg.SaveKey != MessageDefinition.UNSET and runMobileCodeMsg.SaveKey or None),
self.__stateContext.cookie)
response = RunMobileCodeAck(Cookie = self.__stateContext.cookie,
MobileCodeAccepted = success,
Message = msg)
self.writeMessageAndClose(response)
elif signal == CheckMobileCodeResult:
if self.__stateContext.encryptedResult != None:
self.__fsm.signal(self.SIGNAL_CODE_EXECUTION_COMPLETE, None)
return
response = RunMobileCodeAck(Cookie=self.__curState.cookie,
MobileCodeAccepted=True,
Message="Still running")
return 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 __generateEncryptedResultMessage(self):
response = EncryptedMobileCodeResult(Cookie=self.__stateContext.cookie,
RunTime=self.__stateContext.runtime,
RunMobileCodeHash=self.__stateContext.runMobileCodeHash,
EncryptedMobileCodeResultPacket=self.__stateContext.encryptedResult)
if self.__stateContext.encryptedResult != "":
response.Success = True
else: response.Success = False
return response
def __generateDecryptionKeyMessage(self):
decryptionKey, decryptionIv = self.__factory.getDecryptionData(data.Cookie)
response = ResultDecryptionKey(Cookie=self.__stateContext.cookie,
key=decryptionKey,
iv=decryptionIv)
return response
def __mobileCodeComplete(self, signal, data):
self.__stateContext.restoreStateSignal = self.SIGNAL_RESTORE_PURCHASE
if signal == self.SIGNAL_RESTORE_RUNNING:
self.__fsm.signal(data.__class__, data)
elif signal == self.SIGNAL_CODE_EXECUTION_COMPLETE or signal == CheckMobileCodeResult:
if self.__stateContext.billingRate == 0:
self.__fsm.signal(self.SIGNAL_NO_CHARGE, data)
return
response = self.__generateEncryptedResultMessage()
return self.writeMessageAndClose(response)
def __handleFinished(self, signal, data):
self.__stateContext.restoreStateSignal = self.SIGNAL_RESTORE_FINISHED
if signal == self.SIGNAL_RESTORE_RUNNING:
self.__fsm.signal(data.__class__, data)
elif signal == PurchaseDecryptionKey or signal == self.SIGNAL_NO_CHARGE:
if signal == PurchaseDecryptionKey:
pass
response = self.__generateDecryptionKeyMessage()
self.writeMessageAndClose(response)
elif signal == CheckMobileCodeResult:
if self.__stateContext.paid >= self.__stateContext.billingRate:
response = self.__generateDecryptionKeyMessage()
else:
response = self.__generateEncryptedResultMessage()
self.writeMessageAndClose(response)
def __handleCheckMobileCodeResult(self, prot, msg):
msgObj = msg.data()
self.__curState = self.__factory.getSessionState(msgObj.Cookie)
if not self.__curState:
return self.__error("No such session found for cookie %s." % msgObj.Cookie,
fatal=True)
elif self.__curState.state == self.STATE_RUNNING:
response = MessageData.GetMessageBuilder(RunMobileCodeAck)
response["Cookie"].setData(self.__curState.cookie)
response["MobileCodeAccepted"].setData(True)
response["Message"].setData("Still running")
return self.writeMsgAndClose(response)
elif self.__curState.state == self.STATE_PURCHASE:
# curState.encryptedResult is an already serialized packet.
# so we can't use writeMsg. Have to write, then close
# Can't do this here: self.writeMsgAndClose(self.__curState.encryptedResult)
self.transport.write(self.__curState.encryptedResult)
return self.transport.loseConnection()
if self.__curState.state not in [self.STATE_RUNNING, self.STATE_PURCHASE]:
return self.__error("Invalid command. Cannot check result in state (%s) cookie %s" %
(self.__curState.state, self.__curState.cookie),
fatal=False)
def __handlePurchase(self, prot, msg):
msgObj = msg.data()
self.__curState = self.__factory.getSessionState(msgObj.Cookie)
if self.__curState.state != self.STATE_PURCHASE:
return self.__error("Invalid command. Not in correct state for purchase (%s)" % self.__curState.state,
fatal=False)
if not self.__factory.validatePurchase(msgObj.Cookie,
msgObj.Receipt, msgObj.ReceiptSignature):
return self.__error("Invalid purchase receipt", fatal=True)
decryptionKey, decryptionIv = self.__factory.getDecryptionData(msgObj.Cookie)
if not decryptionKey or not decryptionIv:
return self.__error("Unexpected failure in getDecryptionData!", fatal=True)
response = MessageData.GetMessageBuilder(ResultDecryptionKey)
response["Cookie"].setData(msgObj.Cookie)
response["key"].setData(decryptionKey)
response["iv"].setData(decryptionIv)
self.__state = self.STATE_FINISHED
packetTrace(logger, response, "%s sending key %s, iv %s" % (msgObj.Cookie,
binascii.hexlify(decryptionKey),
binascii.hexlify(decryptionIv)))
self.writeMsgAndClose(response)
def __handleRerequest(self, prot, msg):
msgObj = msg.data()
self.__curState = self.__factory.getSessionState(msgObj.Cookie)
if not self.__curState.state == self.STATE_FINISHED:
return self.__error("Cannot re-request a key until the session is finished", fatal=False)
self.__state = self.STATE_REREQUEST
msgObj = msg.data()
decryptionKey, decryptionIv = self.__factory.getDecryptionData(msgObj.Cookie)
if not decryptionKey or not decryptionIv:
return self.__error("No decryption key found", fatal=False)
response = MessageData.GetMessageBuilder(ResultDecryptionKey)
response["Cookie"].setData(msgObj.Cookie)
response["key"].setData(decryptionKey)
response["iv"].setData(decryptionIv)
packetTrace(logger, response, "%s re-sending key %s, iv %s" % (msgObj.Cookie,
binascii.hexlify(decryptionKey),
binascii.hexlify(decryptionIv)))
self.writeMsgAndClose(response)
def __init__(self):
print "@RIP_Client.__init__()"
self.nonce = os.urandom(8).encode('hex')
#self.nonce = randint(1000, 9999)
self.storage = MessageStorage()
self.STATES = [CLOSED(), SYN_SENT(), ESTAB()]
class RipProtocol(StackingProtocolMixin, Protocol):
def __init__(self):
self.buffer = ""
self.storage = MessageStorage()
self.TCB = TransmissionControlBlock()
self.higherTransport = RipTransport(self.transport, self, self.TCB)
self.fsm = FSM("RipStateMachine")
self.fsm.addState("Closed", ("Send_SNN", "SNN-SENT"))
self.fsm.addState("Listening", ("SNN_Received", "ACK-Sent"))
self.fsm.addState("SNN-SENT", ("SNN_ACK_Received", "Established"),
onEnter=self.sendSyn)
self.fsm.addState("ACK-Sent", ("ACK_Received", "Established"),
onEnter=self.sendSynAck)
self.fsm.addState("Established", ("CLOSE_Requested", "CLOSE-REQ"),
("CLOSE_Received", "CLOSE-RECV"),
onEnter=self.sendAck)
self.fsm.addState("CLOSE-REQ", ("CLOSE_ACK_Received", "Closed"))
self.fsm.addState("CLOSE-RECV", ("CLOSE_ACK_Sent", "Listening"),
onEnter=self.sendFinAck)
def connectionMade(self):
self.higherTransport = RipTransport(self.transport, self, self.TCB)
if self.factory.State == "CLIENT": self.fsm.start("Closed")
elif self.factory.State == "SERVER": self.fsm.start("Listening")
if (self.factory.State == "CLIENT"
and self.fsm.currentState() == "Closed"):
self.fsm.signal("Send_SNN", '')
def dataReceived(self, data):
self.buffer += data
self.storage.update(data)
for msg in self.storage.iterateMessages():
try:
message = msg
except Exception, e:
print "Error", e
return
# If SNN is received
if (message.sequence_number_notification_flag
and self.fsm.currentState() == "Listening"):
print "%s SNN Received with Sequence No %s" % (
self.factory.State, message.sequence_number)
self.fsm.signal("SNN_Received", message)
# If SNN-ACK in received
elif (message.sequence_number_notification_flag
and message.acknowledgement_flag and
message.acknowledgement_number == self.TCB.sequence_number
and self.fsm.currentState() == "SNN-SENT"):
print "%s: SNN-ACK Received with Sequence No %s" % (
self.factory.State, message.sequence_number)
self.fsm.signal("SNN_ACK_Received", message)
self.makeHigherConnection(self.higherTransport)
# If final ACK is received
elif (not message.sequence_number_notification_flag
and message.acknowledgement_flag and
message.acknowledgement_number == self.TCB.sequence_number
and self.fsm.currentState() == "ACK-Sent"):
if self.checkMessageIntegrity(message):
peerCert = message.certificate
signedNonce = peerCert[0]
cert = self.TCB.peerCert
caCert = self.TCB.peerCaCert
if (self.verifySignedNonce(cert, caCert, signedNonce)):
self.fsm.signal("ACK_Received", message)
print "%s: ACK Received with Sequence No %s" % (
self.factory.State, message.sequence_number)
self.TCB.next_seq_expected = message.sequence_number + 1
self.makeHigherConnection(self.higherTransport)
# If a message is received by one of the users
elif (message.data != ""
and self.fsm.currentState() == "Established"
and message.sessionID
== str(self.TCB.peer_nonce) + str(self.TCB.nonce)
and not message.acknowledgement_flag
and not message.close_flag
and message.sequence_number == self.TCB.next_seq_expected):
if self.checkMessageIntegrity(message):
self.TCB.next_seq_expected = message.sequence_number + len(
str(message.data))
# Add the message to the list of received messages
self.TCB.receivedMessages[
message.sequence_number] = message
self.sendMessageAck(message.data)
# If ACK for the packet is received
elif (self.fsm.currentState() == "Established"
and message.sessionID
== str(self.TCB.peer_nonce) + str(self.TCB.nonce)
and message.acknowledgement_flag and not message.close_flag):
# Remove the acknowledged message from the retransmission buffer
for key, message in list(
self.TCB.retransmissionBuffer.items()):
if key <= message.acknowledgement_number:
del self.TCB.retransmissionBuffer[key]
# If CLOSING ACK is received
elif (message.close_flag and message.acknowledgement_number
and self.fsm.currentState() == "CLOSE-REQ"):
print "%s: CLOSE ACK Received with Sequence Number %s" % (
self.factory.State, message.sequence_number)
if self.TCB.retransmissionBuffer:
self.retransmitPackets()
self.fsm.signal("CLOSE_ACK_Received", message)
self.closeConnection()
elif (message.sequence_number_notification_flag
and self.fsm.currentState() == "Established"):
self.closeConnection()
def __init__(self):
print "@RIP_Server.__init__()"
self.nonce = os.urandom(8).encode('hex')
self.storage = MessageStorage()
self.STATES = [LISTEN(self.nonce), SYN_RECV(), ESTAB()]
def __init__(self):
self.__storage = MessageStorage(ReprogrammingResponse)
self.__requests = {}
self.__reqId = 0
class GateProtocol(Protocol):
@classmethod
def GetNextResvId(cls):
return random.randint(0, 2**32)
def __init__(self):
self.__buffer = MessageStorage()
self.__listenLookup = {}
self.__connLookup = {}
def reservePort(self, srcPort, callbackPort, serverEndpoint):
resvId = self.GetNextResvId()
g_logger.info(
"Sending G2G port reservation to gate with ID %d for srcPort %d" %
(resvId, srcPort))
resv = Gate2GateReservation(
resvType=Gate2GateReservation.RESV_TYPE_LISTEN,
resvId=resvId,
callbackAddr=self.transport.getHost().host,
callbackPort=callbackPort,
srcPort=srcPort)
self.transport.write(resv.__serialize__())
self.__listenLookup[resvId] = serverEndpoint
def connect(self, dstAddr, dstPort, callbackPort, clientEndpoint):
resvId = self.GetNextResvId()
g_logger.info(
"Sending G2G outbound reservation to gate with ID %d for connection to %s:%d"
% (resvId, dstAddr, dstPort))
resv = Gate2GateReservation(
resvType=Gate2GateReservation.RESV_TYPE_CONNECT,
resvId=resvId,
callbackAddr=self.transport.getHost().host,
callbackPort=callbackPort,
dstAddr=dstAddr,
dstPort=dstPort)
self.transport.write(resv.__serialize__())
self.__connLookup[resvId] = clientEndpoint
def dataReceived(self, data):
self.__buffer.update(data)
for g2gMessage in self.__buffer.iterateMessages():
self.__handleG2gResponse(g2gMessage)
def __handleG2gResponse(self, g2gMessage):
if g2gMessage.resvType == Gate2GateReservation.RESV_TYPE_LISTEN:
self.__handleG2gResvListen(g2gMessage)
elif g2gMessage.resvType == Gate2GateReservation.RESV_TYPE_CONNECT:
self.__handleG2gResvConnect(g2gMessage)
def __handleG2gResvListen(self, g2gMessage):
resvId, respType = g2gMessage.resvId, g2gMessage.respType
error = ""
if not self.__listenLookup.has_key(resvId):
errReporter.warning("Unexpected response with Resevation Id %d" %
resvId)
return
if g2gMessage.success:
if respType == Gate2GateResponse.RESP_TYPE_INITIAL:
# this is a generic reservation response
self.__listenLookup[resvId].gateback(
(GateService.CALLBACK_TYPE_RESERVATION, g2gMessage.srcAddr,
g2gMessage.srcPort, g2gMessage.msg))
elif respType == Gate2GateResponse.RESP_TYPE_CALLBACK:
# this means there's a newly spawned connection on the reserved port
self.__listenLookup[resvId].gateback(
(GateService.CALLBACK_TYPE_SPAWN, g2gMessage.dstAddr,
g2gMessage.dstPort, g2gMessage.connPort, g2gMessage.msg))
else:
error = "Got an unexpected response type %s" % respType
else:
error = "Gate reported error msg: %s" % g2gMessage.msg
if error:
self.__listenLookup[resvId].gateerr(Failure(error))
del self.__listenLookup[resvId]
def __handleG2gResvConnect(self, g2gMessage):
resvId, respType = g2gMessage.resvId, g2gMessage.respType
error = ""
if not self.__connLookup.has_key(resvId):
errReporter.warning("Unexpected response for reservation Id %d" %
resvId)
return
if g2gMessage.success:
if respType == Gate2GateResponse.RESP_TYPE_INITIAL:
self.__connLookup[resvId].gateback(
(GateService.CALLBACK_TYPE_CONNECT, g2gMessage.srcAddr,
g2gMessage.srcPort, g2gMessage.dstAddr,
g2gMessage.dstPort, g2gMessage.msg))
elif respType == Gate2GateResponse.RESP_TYPE_CALLBACK:
self.__connLookup[resvId].gateback(
(GateService.CALLBACK_TYPE_CONNECT_COMPLETE,
g2gMessage.connPort, g2gMessage.msg))
else:
error = "Got an unexpected response type %s" % respType
else:
error = "Gate reported error msg: %s" % g2gMessage.msg
if error:
self.__listenLookup[resvId].gateerr(Failure(error))
del self.__listenLookup[resvId]
