Skip to content

Commit

Permalink
Add query param for not encoding payload
Browse files Browse the repository at this point in the history
  • Loading branch information
quentin on chickenita committed Jun 13, 2024
1 parent 61a4edb commit b284b5f
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 8 deletions.
14 changes: 11 additions & 3 deletions pypeman/message.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,16 +88,24 @@ def add_context(self, key, msg):
payload=copy.deepcopy(msg.payload),
)

def to_dict(self):
def to_dict(self, encode_payload=True):
"""
Convert the current message object to a dict. Payload is pickled.
Convert the current message object to a dict.
Payload is pickled and b64 encoded if encode_payload not set to False
:return: A dict with an equivalent of message
"""
result = {}
result['timestamp'] = self.timestamp.strftime(DATE_FORMAT)
result['uuid'] = self.uuid
result['payload'] = base64.b64encode(pickle.dumps(self.payload)).decode('ascii')
if encode_payload:
result['payload'] = base64.b64encode(pickle.dumps(self.payload)).decode('ascii')
else:
try:
result['payload'] = str(self.payload)
except Exception:
default_logger.warning("Cannot convert to string payload %r, pickling it")
result['payload'] = base64.b64encode(pickle.dumps(self.payload)).decode('ascii')
result['meta'] = self.meta
result['ctx'] = {}

Expand Down
31 changes: 26 additions & 5 deletions pypeman/plugins/remoteadmin/views.py
Original file line number Diff line number Diff line change
Expand Up @@ -133,13 +133,18 @@ async def replay_msg(request, ws=None):
:params channel: The channel name.
:params msg_id: The message id to replay.
:queryparams:
encode_payload (Bool, default=False): Force pickling and encoding the payload or not
"""
channelname = request.match_info['channelname']
message_id = request.match_info['message_id']
args = request.rel_url.query
encode_payload = args.get("encode_payload", False)
chan = get_channel(channelname)
try:
msg_res = await chan.replay(message_id)
result = msg_res.to_dict()
result = msg_res.to_dict(encode_payload=encode_payload)
except IndexError:
message = f"Cannot replay msg, id {message_id} probably doesn't exists"
logger.error(message)
Expand All @@ -160,16 +165,22 @@ async def view_msg(request, ws=None):
:params channelname: The channel name.
:params message_id: The message id to view
:queryparams:
encode_payload (Bool, default=False): Force pickling and encoding the payload or not
"""

channelname = request.match_info['channelname']
message_id = request.match_info['message_id']

args = request.rel_url.query
encode_payload = args.get("encode_payload", False)

chan = get_channel(channelname)

try:
msg_res = await chan.message_store.get_msg_content(message_id)
result = msg_res.to_dict()
result = msg_res.to_dict(encode_payload=encode_payload)
except IndexError:
message = f"Cannot view msg, id {message_id} probably doesn't exists"
logger.error(message)
Expand All @@ -190,14 +201,20 @@ async def preview_msg(request, ws=None):
:params channelname: The channel name.
:params message_id: The message id to preview
:queryparams:
encode_payload (Bool, default=False): Force pickling and encoding the payload or not
"""
channelname = request.match_info['channelname']
message_id = request.match_info['message_id']

args = request.rel_url.query
encode_payload = args.get("encode_payload", False)

chan = get_channel(channelname)
try:
msg_res = await chan.message_store.get_preview_str(message_id)
result = msg_res.to_dict()
result = msg_res.to_dict(encode_payload=encode_payload)
except IndexError:
message = f"Cannot preview msg, id {message_id} probably doesn't exists"
logger.error(message)
Expand Down Expand Up @@ -248,11 +265,15 @@ async def backport_old_client(request):
elif cmd_method == "preview_msg":
message_id = params[1]
request.match_info["message_id"] = message_id
await preview_msg(request, ws=ws)
query_url = request.rel_url.with_query({"encode_payload": True})
new_req = request.clone(rel_url=query_url)
await preview_msg(request=new_req, ws=ws)
elif cmd_method == "view_msg":
message_id = params[1]
request.match_info["message_id"] = message_id
await view_msg(request=request, ws=ws)
query_url = request.rel_url.with_query({"encode_payload": True})
new_req = request.clone(rel_url=query_url)
await view_msg(request=new_req, ws=ws)
elif cmd_method == "replay_msg":
message_id = params[1]
request.match_info["message_id"] = message_id
Expand Down

0 comments on commit b284b5f

Please sign in to comment.