def on_msg_handle(self, callback: Callable[[Dict[str, Any]], None],
body: Dict[str, Any], message: Message) -> None:
"""自定义的方法, 会自动预处理数据, 和报警, 如果是成功或者是skip错误, 就会自动ack"""
consumer_name: str = callback.__name__
tag: Optional[str] = None
# 如果消息带有id,则直接使用消息的id
if self.msg_id_key:
tag = body.get(self.msg_id_key, None)
if not tag:
message_data_list = [
message.delivery_info.get(key, None)
for key in ["exchange", "routing_key", "delivery_tag"]
]
tag = ".".join([i for i in message_data_list if i])
log.info(
f"{self.worker_name}:{consumer_name} recv msg, delivery info:{message.delivery_info} body:{body}"
)
try:
mq_context.set(f"{consumer_name}:{tag}")
callback(body)
message.ack()
except SkipException as e:
message.ack()
log.info(f"recv skip error:{e} auto ack")
except Exception as e:
msg: str = f"{consumer_name} failed to handle recv msg, error: {e}"
log.exception(msg)
def handle(self, message: kombu.Message) -> None:
context = self.baseplate.make_context_object()
try:
# We place the call to ``baseplate.make_server_span`` inside the
# try/except block because we still want Baseplate to see and
# handle the error (publish it to error reporting)
with self.baseplate.make_server_span(context, self.name) as span:
delivery_info = message.delivery_info
message_body = None
message_body = message.decode()
span.set_tag("kind", "consumer")
span.set_tag("amqp.routing_key",
delivery_info.get("routing_key", ""))
span.set_tag("amqp.consumer_tag",
delivery_info.get("consumer_tag", ""))
span.set_tag("amqp.delivery_tag",
delivery_info.get("delivery_tag", ""))
span.set_tag("amqp.exchange",
delivery_info.get("exchange", ""))
self.handler_fn(context, message_body, message)
except Exception as exc:
logger.exception(
"Unhandled error while trying to process a message. The message "
"has been returned to the queue broker.")
if self.error_handler_fn:
self.error_handler_fn(context, message_body, message, exc)
else:
message.requeue()
if isinstance(exc, FatalMessageHandlerError):
logger.info("Recieved a fatal error, terminating the server.")
raise
else:
message.ack()
def on_message(self, body, message: Message):
try:
with transaction.atomic(using='default'):
process_email_notification(body)
message.ack()
except Exception:
# _log.error(f'cant save email notification {body}')
message.reject()
def _process_message(self, body, message: Message):
# ack message first, metrics don't have to be exact
message.ack()
if body['type'] == 'flag_submit':
self.flag_submits_metric.labels(**body['data']).inc(
amount=body.get('value', 1), )
else:
self.app.logger.error('Unknown metric type in %s', body)
def respond(self, message: Message):
body = message.body
properties = message.properties
print(f"Received message: {body}")
reply_to = properties.get('reply_to')
if reply_to:
response = f"Reply: {body}"
correlation_id = properties.get('correlation_id')
print(f"Replying to {reply_to}")
self.producer.publish(response, reply_to, correlation_id=correlation_id)
message.ack()
def _on_message(self, body: Any, message: kombu.Message) -> None:
"""
Default option for a consumer callback, prints out message and message data.
:param body: contains the body of the message sent
:param message: contains meta data about the message sent (ie. delivery_info)
"""
if self._debug:
print(f"Message Received @ {datetime.now()}", flush=True)
message.ack()
for fun in self._callbacks:
fun(body, message)
def test_realtime_loader_invalid_protobuf(setup_database):
message = Message()
rt_loader = RTReloader("hostname", "queue", "exchange", 10)
with app.app_context():
resp_rt = rt_loader._on_request(message)
assert len(resp_rt) == 2
assert "encoding without a string argument" == resp_rt["reason"]
assert resp_rt["duration"] > 0
def _experiment_message_handler(self, body: str, msg: Message):
try:
exp = ExperimentMessage.from_dict(json.loads(body))
self.logger.info(f'Received experiment message for process {exp.process_uuid} (index {exp.experiment_index} for submission {exp.submission_uuid})')
self.terra_exporter.export(exp.process_uuid, exp.submission_uuid, exp.job_id)
self.logger.info(f'Exported experiment for process uuid {exp.process_uuid} (--index {exp.experiment_index} --total {exp.total} --submission {exp.submission_uuid})')
self.log_complete_assay(exp.job_id, exp.process_id)
self.producer.publish(json.loads(body),
exchange=self.publish_queue_config.exchange,
routing_key=self.publish_queue_config.routing_key,
retry=self.publish_queue_config.retry,
retry_policy=self.publish_queue_config.retry_policy)
msg.ack()
except Exception as e:
self.logger.error(f'Failed to export experiment message with body: {body}')
self.logger.exception(e)
def test_realtime_loader_valid_protobuf(setup_database):
task = task_pb2.Task()
task.action = task_pb2.LOAD_REALTIME
task.load_realtime.queue_name = "kirin_load_realtime"
task.load_realtime.contributors.append(COTS_CONTRIBUTOR_ID)
task.load_realtime.begin_date = "2015-11-04"
task.load_realtime.end_date = "2015-11-04"
message = Message()
message.body = task.SerializeToString()
rt_loader = RTReloader("hostname", "queue", "exchange", 10)
with app.app_context():
resp_rt = rt_loader._on_request(message)
assert len(resp_rt) == 6
assert [COTS_CONTRIBUTOR_ID] == resp_rt["contributors"]
assert task.load_realtime.queue_name == resp_rt["routing_key"]
assert 0 == resp_rt["output_trip_update_count"]
assert 46 == resp_rt["output_feed_size"]
# currently: error while trying to publish the message
assert "'str' object has no attribute 'clone'" == resp_rt["reason"]
assert resp_rt["duration"] > 0
def on_message(message: Message):
logger.debug("From queue: %s, received: %s", config.amqp_queue,
message.body)
try:
validate_payload(message.body)
except ValueError as err:
logger.error("%s Message: %s", err, message.body)
message.reject()
logger.warning("Rejected: %s", message.body)
return
try:
es.index(index=config.es_index,
body=message.body,
doc_type=config.es_doc_type)
except ElasticsearchException as err:
logger.error("Cannot modify index %s. Error: %s", config.es_index,
err)
message.reject()
logger.warning("Rejected: %s", message.body)
return
logger.debug(
"To index: %s, type: %s, document: %s",
config.es_index,
config.es_doc_type,
message.body,
)
message.ack()
logger.debug("Acked: %s", message.body)
def on_message(message: Message):
logger.debug("From queue: %s, received: %s", config.amqp_queue,
message.body)
try:
email = validate_payload(message.body)
msg = compose_email(email)
except ValueError as err:
logger.error("%s Email Message: %s", err, email)
message.reject()
logger.warning("Rejected: %s", message.body)
return
try:
with SMTP(host=config.smtp_host, port=config.smtp_port) as smtp:
smtp.starttls(
) #TODO: We should check server's certificate here.
smtp.login(config.smtp_user, config.smtp_pass)
smtp.send_message(msg)
except (ConnectionRefusedError, SMTPException) as err:
logger.error("Cannot send email message: %s Message: %s", err,
email)
message.reject()
logger.warning("Rejected: %s", message.body)
return
logger.info("Email sent: %s", email)
message.ack()
logger.debug("Acked: %s", message.body)
def analyze(self, body: Dict, message: Message):
"""
Analyze incoming Tweet and extract indicators of compromise.
:param body: Json message
:param message: AMPQ message object
"""
try:
tweet: Tweet = Tweet.parse_obj(body)
if tweet.retweeted_status is not None:
tweet: Tweet = Tweet.parse_obj(tweet.retweeted_status)
elif tweet.quoted_status is not None:
tweet: Tweet = Tweet.parse_obj(tweet.quoted_status)
if hasattr(tweet.extended_tweet, 'full_text'):
text = tweet.extended_tweet.full_text
entities = tweet.extended_entities
else:
text = tweet.text
entities = tweet.entities
if entities is not None:
for url in entities.urls:
text = text.replace(url.url, url.expanded_url)
sha256s, md5s = self.extract_hashes(text)
malware: bool = False
for h in set(sha256s + md5s):
h = h.lower()
for i in range(0, 6):
if self.redis.getbit(h, i):
malware = True
break
if malware:
break
if malware:
if not self.redis.getbit(tweet.id_str, 0):
tweet_url: str = "https://twitter.com/" + tweet.user.screen_name + "/status/" + tweet.id_str
logging.info(tweet_url)
self.redis.setbit(tweet.id_str, 0, 1)
self.redis.expire(tweet.id_str, 172800) # 48 hours
self.send_dashboard({
"created_at": tweet.created_at,
"id": tweet.id_str,
"text": tweet.text,
"lang": tweet.lang,
"name": tweet.user.name,
"screen_name": tweet.user.screen_name,
"profile_image_url_https": tweet.user.profile_image_url_https
})
malware_family: str = ""
for hashtag in [hashtag.text.lower() for hashtag in tweet.entities.hashtags]:
for name, keywords in self.malpedia_families.items():
if hashtag in keywords:
malware_family = name
break
if malware_family != "":
self.send_telegram("#" + malware_family.replace('.', '_') + " " + tweet_url)
self.send_threatfox(tweet_url, sha256s, md5s, malware_family)
else:
logging.info("Could not find any malware family in Tweet -> %s" % tweet_url)
self.send_telegram(tweet_url)
mysql = self.engine.connect()
mysql.execute(insert(TableUser).values(
id=tweet.user.id,
name=tweet.user.name,
screen_name=tweet.user.screen_name,
protected=tweet.user.protected,
verified=tweet.user.verified,
created_at=tweet.user.created_at,
profile_image_url_https=tweet.user.profile_image_url_https
).on_duplicate_key_update(
name=tweet.user.name,
screen_name=tweet.user.screen_name,
protected=tweet.user.protected,
verified=tweet.user.verified,
profile_image_url_https=tweet.user.profile_image_url_https))
mysql.execute(insert(TableTweet).values(
id=tweet.id,
user_id=tweet.user.id,
created_at=tweet.created_at,
text=text,
source=tweet.source,
truncated=tweet.truncated,
in_reply_to_user_id=tweet.in_reply_to_user_id,
in_reply_to_status_id=tweet.in_reply_to_status_id,
in_reply_to_screen_name=tweet.in_reply_to_screen_name,
possibly_sensitive=tweet.possibly_sensitive,
lang=tweet.lang
).on_duplicate_key_update(id=tweet.id))
mysql.close()
except pydantic.ValidationError as e:
logging.warning("Failed to decode Tweet ->\n%s" % e)
except pymysql.err.IntegrityError as e:
logging.warning("Tweet already in table ->\n%s" % e)
except exc.SQLAlchemyError as e:
logging.warning("Failed to insert Tweet into database ->\n%s" % e)
finally:
message.ack()
def on_task(self, body, message: Message):
self.pages_collection.insert_one(body)
message.ack()
def on_message(self, body: AnyStr, message: Message) -> None:
print('Got message: {0}'.format(body))
message.ack()