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))
Example #3
0
    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 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 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()