Sign callback data to prevent tampering with it

Author: pietroalbini

botogram/__init__.py

@@ -34,7 +34,6 @@
 from .runner import run
 from .objects import *
 from .utils import usernames_in
-from .callbacks import buttons
 
 
 # This code will simulate the Windows' multiprocessing behavior if the

botogram/api.py

@@ -182,3 +182,7 @@ def file_content(self, path):
         response = requests.get(url)
 
         return response.content
+
+    @property
+    def token(self):
+        return self._api_key

botogram/callbacks.py

@@ -17,25 +17,34 @@
 #   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 #   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 
+import base64
+import binascii
+import hashlib
+
+from . import crypto
+
+
+DIGEST = hashlib.md5
+DIGEST_LEN = 16
+
 
 class ButtonsRow:
     """A row of an inline keyboard"""
 
-    def __init__(self):
+    def __init__(self, bot):
         self._content = []
+        self._bot = bot
 
     def url(self, label, url):
         """Open an URL when the button is pressed"""
         self._content.append({"text": label, "url": url})
 
     def callback(self, label, callback, data=None):
         """Trigger a callback when the button is pressed"""
-        if data is not None:
-            msg = "%s\0%s" % (callback, data)
-        else:
-            msg = callback
-
-        self._content.append({"text": label, "callback_data": msg})
+        self._content.append({
+            "text": label,
+            "callback_data": get_callback_data(self._bot, callback, data),
+        })
 
     def switch_inline_query(self, label, query="", current_chat=False):
         """Switch the user to this bot's inline query"""
@@ -54,12 +63,13 @@ def switch_inline_query(self, label, query="", current_chat=False):
 class Buttons:
     """Factory for inline keyboards"""
 
-    def __init__(self):
+    def __init__(self, bot):
         self._rows = {}
+        self._bot = bot
 
     def __getitem__(self, index):
         if index not in self._rows:
-            self._rows[index] = ButtonsRow()
+            self._rows[index] = ButtonsRow(self._bot)
         return self._rows[index]
 
     def _serialize_attachment(self):
@@ -72,27 +82,71 @@ def _serialize_attachment(self):
         return {"inline_keyboard": rows}
 
 
-def buttons():
-    """Create a new inline keyboard"""
-    return Buttons()
+def parse_callback_data(bot, raw):
+    """Parse the callback data generated by botogram and return it"""
+    raw = raw.encode("utf-8")
 
+    if len(raw) < 32:
+        raise crypto.TamperedMessageError
 
-def parse_callback_data(data):
-    """Parse the callback data generated by botogram and return it"""
-    if "\0" in data:
-        name, custom = data.split("\0", 1)
-        return name, custom
-    else:
-        return data, None
+    try:
+        prelude = base64.b64decode(raw[:32])
+    except binascii.Error:
+        raise crypto.TamperedMessageError
+
+    signature = prelude[:16]
+    name = prelude[16:]
+    data = raw[32:]
+
+    if not crypto.compare(crypto.get_hmac(bot, name + data), signature):
+        raise crypto.TamperedMessageError
+
+    return name, data.decode("utf-8")
+
+
+def get_callback_data(bot, name, data=None):
+    """Get the callback data for the provided name and data"""
+    name = hashed_callback_name(name)
+
+    if data is None:
+        data = ""
+    data = data.encode("utf-8")
+
+    if len(data) > 32:
+        raise ValueError(
+            "The provided data is too big (%s bytes), try to reduce it to "
+            "32 bytes" % len(data)
+        )
+
+    # Get the signature of the hook name and data
+    signature = crypto.get_hmac(bot, name + data)
+
+    # Base64 the signature and the hook name together to save space
+    return (base64.b64encode(signature + name) + data).decode("utf-8")
+
+
+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
+    return DIGEST(name.encode("utf-8")).digest()[:8]
 
 
 def process(bot, chains, update):
     """Process a callback sent to the bot"""
+    try:
+        name, data = parse_callback_data(bot, update.callback_query._data)
+    except crypto.TamperedMessageError:
+        bot.logger.warn(
+            "The user tampered with the #%s update's data. Skipped it."
+            % update.update_id
+        )
+        return
+
     for hook in chains["callbacks"]:
         bot.logger.debug("Processing update #%s with the hook %s" %
                          (update.update_id, hook.name))
 
-        result = hook.call(bot, update)
+        result = hook.call(bot, update, name, data)
         if result is True:
             bot.logger.debug("Update #%s was just processed by the %s hook" %
                              (update.update_id, hook.name))

botogram/crypto.py

@@ -0,0 +1,67 @@
+# Copyright (c) 2015-2017 The Botogram Authors (see AUTHORS)
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+#   The above copyright notice and this permission notice shall be included in
+#   all copies or substantial portions of the Software.
+#
+#   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+#   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+#   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+#   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+#   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+#   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+#   DEALINGS IN THE SOFTWARE.
+
+import hmac
+
+
+DIGEST = "md5"
+DIGEST_LEN = 16
+
+
+class TamperedMessageError(Exception):
+    pass
+
+
+def generate_secret_key(bot):
+    """Generate the secret key of the bot"""
+    mac = hmac.new(bot.api.token.encode("utf-8"), digestmod=DIGEST)
+    mac.update(b"botogram" + bot.itself.username.encode("utf-8"))
+    return mac.digest()
+
+
+def get_hmac(bot, data):
+    """Get the HMAC of a piece of data"""
+    mac = hmac.new(generate_secret_key(bot), digestmod=DIGEST)
+    mac.update(data)
+    return mac.digest()
+
+
+def sign_data(bot, data):
+    """Return a signed version of the data, to prevent tampering with it"""
+    return get_hmac(bot, data) + data
+
+
+def verify_signature(bot, untrusted):
+    """Check if the untrusted data is correctly signed, and return it"""
+    if len(untrusted) < DIGEST_LEN:
+        raise TamperedMessageError
+
+    signature = untrusted[:DIGEST_LEN]
+    data = untrusted[DIGEST_LEN:]
+
+    if not hmac.compare_digest(get_hmac(bot, data), signature):
+        raise TamperedMessageError
+
+    return data
+
+
+def compare(a, b):
+    """Safely compare two values"""
+    return hmac.compare_digest(a, b)

botogram/frozenbot.py

@@ -23,6 +23,7 @@
 from . import utils
 from . import objects
 from . import api as api_module
+from .callbacks import Buttons
 
 
 class FrozenBotError(Exception):
@@ -234,6 +235,10 @@ def register_update_processor(self, kind, processor):
         """Register a new update processor"""
         raise FrozenBotError("Can't register new update processors at runtime")
 
+    def buttons(self):
+        """Create a new inline keyboard"""
+        return Buttons(self)
+
     # This helper manages the translation
 
     def _(self, message, **args):

botogram/hooks.py

@@ -19,7 +19,7 @@
 
 import re
 
-from .callbacks import parse_callback_data
+from .callbacks import hashed_callback_name
 
 
 class Hook:
@@ -212,14 +212,15 @@ class CallbackHook(Hook):
     """Underlying hook for @bot.callback"""
 
     def _after_init(self, args):
-        self._name = "%s:%s" % (self.component.component_name, args["name"])
+        self._name = hashed_callback_name(
+            "%s:%s" % (self.component.component_name, args["name"])
+        )
 
-    def _call(self, bot, update):
+    def call(self, bot, update, name, data):
         if not update.callback_query:
             return
         q = update.callback_query
 
-        name, data = parse_callback_data(q._data)
         if name != self._name:
             return
 

tests/test_callbacks.py

@@ -20,11 +20,12 @@
 
 import json
 
-from botogram.callbacks import Buttons, parse_callback_data
+from botogram.callbacks import Buttons, parse_callback_data, get_callback_data
+from botogram.callbacks import hashed_callback_name
 
 
-def test_buttons():
-    buttons = Buttons()
+def test_buttons(bot):
+    buttons = Buttons(bot)
     buttons[0].url("test 1", "http://example.com")
     buttons[0].callback("test 2", "test_callback")
     buttons[3].callback("test 3", "another_callback", "data")
@@ -35,20 +36,30 @@ def test_buttons():
         "inline_keyboard": [
             [
                 {"text": "test 1", "url": "http://example.com"},
-                {"text": "test 2", "callback_data": "test_callback"},
+                {
+                    "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"},
             ],
             [
-                {"text": "test 3", "callback_data": "another_callback\0data"},
+                {
+                    "text": "test 3",
+                    "callback_data": get_callback_data(
+                        bot, "another_callback", "data"
+                    ),
+                },
             ],
         ],
     }
 
 
-def test_parse_callback_data():
-    assert parse_callback_data("test") == ("test", None)
-    assert parse_callback_data("test:something") == ("test:something", None)
-    assert parse_callback_data("test\0wow") == ("test", "wow")
+def test_parse_callback_data(bot):
+    raw = get_callback_data(bot, "test_callback", "this is some data!")
+    assert parse_callback_data(bot, raw) == (
+        hashed_callback_name("test_callback"),
+        "this is some data!",
+    )

tests/test_crypto.py

@@ -0,0 +1,45 @@
+# Copyright (c) 2015-2017 The Botogram Authors (see AUTHORS)
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+#   The above copyright notice and this permission notice shall be included in
+#   all copies or substantial portions of the Software.
+#
+#   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+#   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+#   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+#   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+#   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+#   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+#   DEALINGS IN THE SOFTWARE.
+
+import pytest
+
+from botogram.crypto import generate_secret_key, get_hmac, sign_data
+from botogram.crypto import TamperedMessageError, verify_signature
+
+
+def test_generate_secret_key(bot):
+    # Check if the generated key for the test bot is correct
+    key = generate_secret_key(bot)
+    assert key == b'\\\x93aA\xe8\x8d\x9aL\x8c\xfd\x81,D\xeaj\xd0'
+
+
+def test_hmac(bot):
+    expect = b'q\x06\x9c\xc1\xfa\xd1n\xe8\xef\x17\xf6\xd7Z\xb0G\x7f'
+    assert get_hmac(bot, b'test data') == expect
+
+    signed = sign_data(bot, b'test string')
+    assert verify_signature(bot, signed) == b'test string'
+
+    signed += b'a'
+    with pytest.raises(TamperedMessageError):
+        verify_signature(bot, signed)
+
+    with pytest.raises(TamperedMessageError):
+        verify_signature(bot, b'a')