Restrict callbacks to the chat they were sent in

Author: pietroalbini

botogram/callbacks.py

@@ -45,7 +45,7 @@ def generate_callback_data():
             c = ctx()
 
             name = "%s:%s" % (c.component_name(), callback)
-            return get_callback_data(c.bot, name, data)
+            return get_callback_data(c.bot, c.chat(), name, data)
 
         self._content.append({
             "text": label,
@@ -105,7 +105,7 @@ def buttons():
     return Buttons()
 
 
-def parse_callback_data(bot, raw):
+def parse_callback_data(bot, chat, raw):
     """Parse the callback data generated by botogram and return it"""
     raw = raw.encode("utf-8")
 
@@ -121,7 +121,8 @@ def parse_callback_data(bot, raw):
     name = prelude[16:]
     data = raw[32:]
 
-    if not crypto.compare(crypto.get_hmac(bot, name + data), signature):
+    correct = get_signature(bot, chat, name, data)
+    if not crypto.compare(correct, signature):
         raise crypto.TamperedMessageError
 
     if data:
@@ -130,7 +131,7 @@ def parse_callback_data(bot, raw):
         return name, None
 
 
-def get_callback_data(bot, name, data=None):
+def get_callback_data(bot, chat, name, data=None):
     """Get the callback data for the provided name and data"""
     name = hashed_callback_name(name)
 
@@ -145,12 +146,18 @@ def get_callback_data(bot, name, data=None):
         )
 
     # Get the signature of the hook name and data
-    signature = crypto.get_hmac(bot, name + data)
+    signature = get_signature(bot, chat, name, data)
 
     # Base64 the signature and the hook name together to save space
     return (base64.b64encode(signature + name) + data).decode("utf-8")
 
 
+def get_signature(bot, chat, name, data):
+    """Generate a signature for the provided information"""
+    chat_id = str(chat.id).encode("utf-8")
+    return crypto.get_hmac(bot, name + b'\0' + chat_id + b'\0' + data)
+
+
 def hashed_callback_name(name):
     """Get the hashed name of a callback"""
     # Get only the first 8 bytes of the hash to fit it into the payload
@@ -159,8 +166,11 @@ def hashed_callback_name(name):
 
 def process(bot, chains, update):
     """Process a callback sent to the bot"""
+    chat = update.callback_query.message.chat
+    raw = update.callback_query._data
+
     try:
-        name, data = parse_callback_data(bot, update.callback_query._data)
+        name, data = parse_callback_data(bot, chat, raw)
     except crypto.TamperedMessageError:
         bot.logger.warn(
             "The user tampered with the #%s update's data. Skipped it."

tests/test_callbacks.py

@@ -22,50 +22,65 @@
 
 from botogram.callbacks import Buttons, parse_callback_data, get_callback_data
 from botogram.callbacks import hashed_callback_name
+from botogram.components import Component
+from botogram.context import Context
+from botogram.hooks import Hook
 
 
-def test_buttons(bot):
-    buttons = Buttons(bot)
+def test_buttons(bot, sample_update):
+    component = Component("test")
+    hook = Hook(lambda: None, component)
+
+    buttons = Buttons()
     buttons[0].url("test 1", "http://example.com")
     buttons[0].callback("test 2", "test_callback")
     buttons[3].callback("test 3", "another_callback", "data")
     buttons[2].switch_inline_query("test 4")
     buttons[2].switch_inline_query("test 5", "wow", current_chat=True)
 
-    assert buttons._serialize_attachment() == {
-        "inline_keyboard": [
-            [
-                {"text": "test 1", "url": "http://example.com"},
-                {
-                    "text": "test 2",
-                    "callback_data": get_callback_data(bot, "test_callback"),
-                },
-            ],
-            [
-                {"text": "test 4", "switch_inline_query": ""},
-                {"text": "test 5", "switch_inline_query_current_chat": "wow"},
+    with Context(bot, hook, sample_update):
+        assert buttons._serialize_attachment() == {
+            "inline_keyboard": [
+                [
+                    {"text": "test 1", "url": "http://example.com"},
+                    {
+                        "text": "test 2",
+                        "callback_data": get_callback_data(
+                            bot, sample_update.chat(), "test:test_callback",
+                        ),
+                    },
+                ],
+                [
+                    {"text": "test 4", "switch_inline_query": ""},
+                    {
+                        "text": "test 5",
+                        "switch_inline_query_current_chat": "wow"
+                    },
+                ],
+                [
+                    {
+                        "text": "test 3",
+                        "callback_data": get_callback_data(
+                            bot, sample_update.chat(), "test:another_callback",
+                            "data",
+                        ),
+                    },
+                ],
             ],
-            [
-                {
-                    "text": "test 3",
-                    "callback_data": get_callback_data(
-                        bot, "another_callback", "data"
-                    ),
-                },
-            ],
-        ],
-    }
+        }
+
 
+def test_parse_callback_data(bot, sample_update):
+    c = sample_update.chat()
 
-def test_parse_callback_data(bot):
-    raw = get_callback_data(bot, "test_callback", "this is some data!")
-    assert parse_callback_data(bot, raw) == (
+    raw = get_callback_data(bot, c, "test_callback", "this is some data!")
+    assert parse_callback_data(bot, c, raw) == (
         hashed_callback_name("test_callback"),
         "this is some data!",
     )
 
-    raw = get_callback_data(bot, "test_callback")
-    assert parse_callback_data(bot, raw) == (
+    raw = get_callback_data(bot, c, "test_callback")
+    assert parse_callback_data(bot, c, raw) == (
         hashed_callback_name("test_callback"),
         None,
     )