def handle_job_finished(self): data = self._extract_data_from_message((yield self.read_message())) if not Request.is_valid(data, Request.Command.JOB_FINISHED): if Request.is_valid(data, Request.Command.JOB_FAILED): _logger.warn("[%s] Job failed!" % CLIENT) request = Request.request_from_data(data) ConverterWebSocketHandler.notify(NotificationType.JOB_FAILED, request.args) return raise TCPConnectionException("Invalid data given!") request = Request.request_from_data(data) _logger.info("[%s]: Received file transfer notification" % SERVER) job_ID = int(request.args[Request.ARGS_JOB_ID]) file_name = str(job_ID) + CATROBAT_FILE_EXT file_size = int(request.args[Request.ARGS_FILE_SIZE]) file_hash = request.args[Request.ARGS_FILE_HASH] _logger.info("[%s]: Job ID: %d, File size: %d, SHA256: %s" % (CLIENT, job_ID, file_size, file_hash)) if file_size == 0: raise TCPConnectionException("Cannot transfer empty file...", context=request.args) _logger.debug("[%s]: Reply: Accepted!" % SERVER) yield self.send_message( Reply(result=True, msg="Ready for file transfer!")) download_dir = self.server.settings["download_dir"] file_path = os.path.join(download_dir, file_name) max_input_buffer_size = int( self.server.settings["max_input_buffer_size"]) with open(file_path, 'wb+', 0) as f: if file_size <= max_input_buffer_size: f.write((yield self.read_bytes(file_size))) else: transfered_bytes = 0 while transfered_bytes < file_size: buffer_size = min((file_size - transfered_bytes), max_input_buffer_size) f.write((yield self.read_bytes(buffer_size))) transfered_bytes += buffer_size f.seek(0, 0) computed_file_hash = hashlib.sha256(f.read()).hexdigest() if computed_file_hash != file_hash: raise TCPConnectionException( "Given hash value is not equal to computed hash value.", context=request.args) _logger.debug("[%s]: Reply: Accepted!" % SERVER) _logger.info( "OK! Hash is equal to computed hash value. Finished file transfer!" ) yield self.send_message(Reply(result=True, msg="OK! Hash is equal to computed " \ "hash value. Finished file transfer!")) ConverterWebSocketHandler.notify(NotificationType.JOB_FINISHED, request.args)
def receive_greeting(self): data = json.loads((yield self._connection.read_message()).rstrip()) if not Reply.is_valid(data): raise Exception("Invalid reply!") reply = Reply(data[Reply.KEY_RESULT], data[Reply.KEY_MSG]) if not reply.result: _logger.error("[%s]: %s" % (SERVER, reply.msg)) raise Exception("Greeting failed!") _logger.info('[%s]: "%s"' % (SERVER, reply.msg))
def post_processing(self, args): file_path = os.path.join(args["outputDir"], str(args["jobID"]) + CATROBAT_FILE_EXT) if not os.path.isfile(file_path): yield self._connection.send_message( Request( Request.Command.JOB_FAILED, { Request.ARGS_JOB_ID: args["jobID"], Request.ARGS_MSG: "Cannot transfer file! File does not exist!" })) return file_size = os.path.getsize(file_path) with open(file_path, 'rb') as fp: file_hash = hashlib.sha256(fp.read()).hexdigest() args = { Request.ARGS_JOB_ID: args["jobID"], Request.ARGS_FILE_SIZE: file_size, Request.ARGS_FILE_HASH: file_hash } yield self._connection.send_message( Request(Request.Command.JOB_FINISHED, args)) # File transfer ready (reply) data = json.loads((yield self._connection.read_message()).rstrip()) if not Reply.is_valid(data): raise Exception("Invalid reply!") reply = Reply(data[Reply.KEY_RESULT], data[Reply.KEY_MSG]) if not reply.result: _logger.error("[%s]: %s" % (SERVER, reply.msg)) raise Exception("File transfer failed!") _logger.info('[%s]: "%s"' % (SERVER, reply.msg)) fp.seek(0, 0) _logger.info("[%s]: Sending..." % CLIENT) byte_buffer = fp.read(BUFFER_SIZE) while byte_buffer: yield self._connection.send_message(byte_buffer, logging_enabled=False) byte_buffer = fp.read(BUFFER_SIZE) _logger.info("[%s]: Done Sending..." % CLIENT) # File transfer finished (reply) data = json.loads((yield self._connection.read_message()).rstrip()) if not Reply.is_valid(data): raise Exception("Invalid reply!") reply = Reply(data[Reply.KEY_RESULT], data[Reply.KEY_MSG]) if not reply.result: _logger.error("[%s]: %s" % (SERVER, reply.msg)) raise Exception("File transfer failed!") _logger.info('[%s]: "%s"' % (SERVER, reply.msg))
def authenticate(self): _logger.debug('[%s]: Start authentication' % CLIENT) request = Request(Request.Command.AUTH, { Request.ARGS_AUTH_KEY: self._auth_key}) yield self._connection.send_message(request, logging_enabled=False) # do not log key! data = json.loads((yield self._connection.read_message()).rstrip()) if not Reply.is_valid(data): raise Exception("Invalid reply!") reply = Reply(data[Reply.KEY_RESULT], data[Reply.KEY_MSG]) if not reply.result: _logger.error("[%s]: %s" % (SERVER, reply.msg)) raise Exception("Authentication failed!") _logger.info('[%s]: "%s"' % (SERVER, reply.msg))
def send_job_conversion_finished_notification(self, job_ID, exit_code): # Job finished (request) _logger.debug('[%s]: Sending job finished notification' % CLIENT) args = { Request.ARGS_JOB_ID: job_ID, Request.ARGS_RESULT: exit_code, Request.ARGS_MSG: "Job finished" } yield self._connection.send_message(Request(Request.Command.JOB_CONVERSION_FINISHED_NOTIFICATION, args)) # Job finished (reply) data = json.loads((yield self._connection.read_message()).rstrip()) if not Reply.is_valid(data): raise Exception("Invalid reply!") reply = Reply(data[Reply.KEY_RESULT], data[Reply.KEY_MSG]) if not reply.result: _logger.error("[%s]: %s" % (SERVER, reply.msg)) raise Exception("Job finished notification failed!") _logger.info('[%s]: "%s"' % (SERVER, reply.msg))
def handle_authentication(self): data = self._extract_data_from_message((yield self.read_message())) if not Request.is_valid(data, Request.Command.AUTH): raise TCPConnectionException("Invalid data given!") request = Request.request_from_data(data) address = self.connection.address host = address[0] if isinstance(address, tuple) else address allowed_auth_keys = self.server.settings["allowed_auth_keys"] allowed_auth_keys_for_host = [ auth_key["key"] for auth_key in allowed_auth_keys if auth_key["host"] == host ] if len(allowed_auth_keys_for_host) == 0: _logger.warn( "An intruder '%s' might have tried to connect to our server!" % (str(address))) # TODO: block him... # Don't tell the client that this hostname is forbidden raise TCPConnectionException("Invalid AUTH_KEY given.") if request.args[ Request.ARGS_AUTH_KEY] not in allowed_auth_keys_for_host: raise TCPConnectionException("Invalid AUTH_KEY given.") _logger.info("[%s]: Reply: Authentication successful!" % SERVER) yield self.send_message( Reply(result=True, msg="Authentication successful!"))
def post_processing(self, args): file_path = os.path.join(args["outputDir"], str(args["jobID"]) + CATROBAT_FILE_EXT) if not os.path.isfile(file_path): yield self._connection.send_message(Request(Request.Command.JOB_FAILED, { Request.ARGS_JOB_ID: args["jobID"], Request.ARGS_MSG: "Cannot transfer file! File does not exist!" })) return file_size = os.path.getsize(file_path) with open(file_path, 'rb') as fp: file_hash = hashlib.sha256(fp.read()).hexdigest() args = { Request.ARGS_JOB_ID: args["jobID"], Request.ARGS_FILE_SIZE: file_size, Request.ARGS_FILE_HASH: file_hash } yield self._connection.send_message(Request(Request.Command.JOB_FINISHED, args)) # File transfer ready (reply) data = json.loads((yield self._connection.read_message()).rstrip()) if not Reply.is_valid(data): raise Exception("Invalid reply!") reply = Reply(data[Reply.KEY_RESULT], data[Reply.KEY_MSG]) if not reply.result: _logger.error("[%s]: %s" % (SERVER, reply.msg)) raise Exception("File transfer failed!") _logger.info('[%s]: "%s"' % (SERVER, reply.msg)) fp.seek(0, 0) _logger.info("[%s]: Sending..." % CLIENT) byte_buffer = fp.read(BUFFER_SIZE) while byte_buffer: yield self._connection.send_message(byte_buffer, logging_enabled=False) byte_buffer = fp.read(BUFFER_SIZE) _logger.info("[%s]: Done Sending..." % CLIENT) # File transfer finished (reply) data = json.loads((yield self._connection.read_message()).rstrip()) if not Reply.is_valid(data): raise Exception("Invalid reply!") reply = Reply(data[Reply.KEY_RESULT], data[Reply.KEY_MSG]) if not reply.result: _logger.error("[%s]: %s" % (SERVER, reply.msg)) raise Exception("File transfer failed!") _logger.info('[%s]: "%s"' % (SERVER, reply.msg))
def send_job_started_notification(self, job_ID, title, image_URL): _logger.debug('[%s]: Sending job started notification' % CLIENT) args = { Request.ARGS_JOB_ID: job_ID, Request.ARGS_IMAGE_URL: image_URL, Request.ARGS_TITLE: json_encode(title)[1:-1], Request.ARGS_MSG: "Job started" } request = Request(Request.Command.JOB_STARTED_NOTIFICATION, args) yield self._connection.send_message(request) # Job started (reply) data = json.loads((yield self._connection.read_message()).rstrip()) if not Reply.is_valid(data): raise Exception("Invalid reply!") reply = Reply(data[Reply.KEY_RESULT], data[Reply.KEY_MSG]) if not reply.result: _logger.error("[%s]: %s" % (SERVER, reply.msg)) raise Exception("Job started notification failed!") _logger.info('[%s]: "%s"' % (SERVER, reply.msg))
def handle_exception(self, exception, msg): _logger.error("{}!".format(msg)) _logger.exception(exception) yield self.send_message(Reply(result=False, msg="{}! Closing connection.".format(msg))) if isinstance(exception, TCPConnectionException) and exception.context != None: args = exception.context args["msg"] = msg ConverterWebSocketHandler.notify(NotificationType.JOB_FAILED, args) else: ConverterWebSocketHandler.notify(NotificationType.JOB_FAILED, { "msg": msg }) self.connection.print_error_and_close_stream()
def handle_job_started_notification(self): data = self._extract_data_from_message((yield self.read_message())) if not Request.is_valid(data, Request.Command.JOB_STARTED_NOTIFICATION): raise TCPConnectionException("Invalid data given!") request = Request.request_from_data(data) _logger.info("[%s]: Received job start notification" % SERVER) _logger.info("[%s]: Title of Scratch program: '%s'" % (CLIENT, unicode(request.args[Request.ARGS_TITLE]))) _logger.debug("[%s]: Reply: Accepted!" % SERVER) yield self.send_message(Reply(result=True, msg="ACK")) ConverterWebSocketHandler.notify(NotificationType.JOB_STARTED, request.args)
def handle_job_conversion_finished_notification(self, data): if data == None or not Request.is_valid(data, Request.Command.JOB_CONVERSION_FINISHED_NOTIFICATION): raise TCPConnectionException("Invalid data given!") request = Request.request_from_data(data) _logger.info("[%s]: Received job finished notification" % SERVER) _logger.info("[%s]: %s " % (CLIENT, request.args[Request.ARGS_MSG])) exit_code = int(request.args[Request.ARGS_RESULT]) _logger.info("[%s]: Job finished with exit code: %d" % (SERVER, exit_code)) if exit_code != 0: raise TCPConnectionException("Job failed with exit code: %d" % exit_code, context=request.args) _logger.debug("[%s]: Reply: Accepted!" % SERVER) yield self.send_message(Reply(result=True, msg="ACK")) ConverterWebSocketHandler.notify(NotificationType.JOB_CONVERSION_FINISHED, request.args)
def send_job_started_notification(self, job_ID, title, image_URL): _logger.debug('[%s]: Sending job started notification' % CLIENT) args = { Request.ARGS_JOB_ID: job_ID, Request.ARGS_IMAGE_URL: image_URL, Request.ARGS_TITLE: title, Request.ARGS_MSG: "Job started" } request = Request(Request.Command.JOB_STARTED_NOTIFICATION, args) yield self._connection.send_message(request) # Job started (reply) data = json.loads((yield self._connection.read_message()).rstrip()) if not Reply.is_valid(data): raise Exception("Invalid reply!") reply = Reply(data[Reply.KEY_RESULT], data[Reply.KEY_MSG]) if not reply.result: _logger.error("[%s]: %s" % (SERVER, reply.msg)) raise Exception("Job started notification failed!") _logger.info('[%s]: "%s"' % (SERVER, reply.msg))
def run(self): # (I) Send greeting try: yield self.send_message( Reply(result=True, msg="Hello Client, please authenticate!")) except Exception as e: self.handle_exception(e, "Greeting failed") return except: self.handle_exception(sys.exc_info()[0], "Greeting failed") return # (II) Expecting authentication via AuthKey try: yield self.handle_authentication() except Exception as e: self.handle_exception(e, "Authentication failed") return except: self.handle_exception(sys.exc_info()[0], "Authentication failed") return # (III) Expecting job started notification try: yield self.handle_job_started_notification() except Exception as e: self.handle_exception( e, "Processing job started notification failed") return except: self.handle_exception( sys.exc_info()[0], "Processing job started notification failed") return # (IV) Expecting output and/or progress notifications or job finished notification data = None try: while True: data = self._extract_data_from_message((yield self.read_message())) if Request.is_valid( data, Request.Command.JOB_CONVERSION_FINISHED_NOTIFICATION): break elif Request.is_valid(data, Request.Command.JOB_OUTPUT_NOTIFICATION): yield self.handle_job_output_notification(data) elif Request.is_valid( data, Request.Command.JOB_PROGRESS_NOTIFICATION): yield self.handle_job_progress_notification(data) else: raise TCPConnectionException("Invalid data given!") except Exception as e: self.handle_exception(e, "Processing job notification failed") return except: self.handle_exception(sys.exc_info()[0], "Processing job notification failed") return # (V) Expecting job finished notification try: yield self.handle_job_conversion_finished_notification(data) except Exception as e: self.handle_exception( e, "Processing job finished notification failed") return except: self.handle_exception( sys.exc_info()[0], "Processing job finished notification failed") return # (V) Expecting file transfer try: yield self.handle_job_finished() except Exception as e: self.handle_exception(e, "File transfer failed") return except: self.handle_exception(sys.exc_info()[0], "File transfer failed") return _logger.info("Finished! Closing stream.") self.connection.stream.close()