implemented encryption with static key

requirements.txt

@@ -6,5 +6,6 @@ PyJWT==1.7.1
 python-dotenv==0.10.3
 pylint==2.4.2
 simple-chalk==0.1.0
+cryptography==2.7
 pre-commit==1.18.3
 readchar==2.0.1

setup.py

@@ -1,48 +1,49 @@
-#!python
-from setuptools import setup
-
-
-with open("README.md", "r") as fh:
-    long_description = fh.read()
-
-
-def get_requirements():
-    requirements = [
-        'colorama==0.4.1',
-        'paho-mqtt==1.4.0',
-        'peewee==3.11.2',
-        'requests==2.22.0',
-        'PyJWT==1.7.1',
-        'python-dotenv==0.10.3',
-        'simple-chalk==0.1.0',
-        'readchar==2.0.1'
-    ]
-    return requirements
-
-
-setup(
-    name='tmessage',
-    version="0.0.2",
-    author="Haider Ali",
-    author_email="haider.lee23@gmail.com",
-    long_description=long_description,
-    long_description_content_type="text/markdown",
-    url="https://github.com/Haider8/tmessage",
-    license="GPL-3.0",
-    entry_points={
-        'console_scripts': ['tmessage=tmessage.cli:main'],
-    },
-    description="""
-      This is a lightweight and low bandwidth
-      CLI tool which can be used for group
-      communication right from your terminal.
-      """,
-    classifiers=[
-        "Programming Language :: Python :: 3",
-        "Operating System :: OS Independent",
-    ],
-    python_requires='>=3.4',
-    packages=['tmessage'],
-    install_requires=get_requirements(),
-    zip_safe=False,
-)
+#!python
+from setuptools import setup
+
+
+with open("README.md", "r") as fh:
+    long_description = fh.read()
+
+
+def get_requirements():
+    requirements = [
+        'colorama==0.4.1',
+        'paho-mqtt==1.4.0',
+        'peewee==3.11.2',
+        'requests==2.22.0',
+        'PyJWT==1.7.1',
+        'python-dotenv==0.10.3',
+        'simple-chalk==0.1.0',
+        'cryptography==2.7',
+        'readchar==2.0.1'
+    ]
+    return requirements
+
+
+setup(
+    name='tmessage',
+    version="0.0.2",
+    author="Haider Ali",
+    author_email="haider.lee23@gmail.com",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url="https://github.com/Haider8/tmessage",
+    license="GPL-3.0",
+    entry_points={
+        'console_scripts': ['tmessage=tmessage.cli:main'],
+    },
+    description="""
+      This is a lightweight and low bandwidth
+      CLI tool which can be used for group
+      communication right from your terminal.
+      """,
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "Operating System :: OS Independent",
+    ],
+    python_requires='>=3.4',
+    packages=['tmessage'],
+    install_requires=get_requirements(),
+    zip_safe=False,
+)

tmessage/cli.py

@@ -6,10 +6,9 @@
 from simple_chalk import chalk
 import tmessage.auth as auth  # auth.py
 from tmessage.db import store_messages  # db.py
-from tmessage.utils import get_formatted_message  # utils.py
+from tmessage.utils import get_formatted_message, encrypt_message, decrypt_message  # utils.py
 from tmessage.input import Input  # input.py
 
-
 # Initialize colorama
 init()
 
@@ -51,13 +50,14 @@ def on_message(client, userdata, message):
 
     user_name_separator_index = current_msg.find(":") + 1
     user_details = chalk.bgGreen(current_msg[:user_name_separator_index])
-    msg = current_msg[user_name_separator_index + 1:]
+    encrypted_msg = current_msg[user_name_separator_index + 3:-1]
+    msg = decrypt_message(encrypted_msg.encode()).decode()
     if user != CURRENT_USER:
         print('\r\033[K' + user_details + " " + msg)
         print('\r\033[K' + MESSAGE_INPUT.get_buffer(), end='', flush=True)
         _, _, message = current_msg.partition("] ")
         if IS_STORE:
-            store_messages(user, message)
+            store_messages(user, msg)
 
 
 def main():
@@ -89,7 +89,7 @@ def main():
         while True:
             raw_msg = MESSAGE_INPUT.get_input()
             formatted_msg = get_formatted_message(raw_msg)
-            pub_msg = f"[{user_name}] {displayed_name}: {formatted_msg}"
+            pub_msg = f"[{user_name}] {displayed_name}: {encrypt_message(formatted_msg.encode())}"
             if raw_msg != "":
                 MQTT_CLIENT.publish(MQTT_TOPIC, pub_msg)
                 if IS_STORE:

tmessage/utils.py

@@ -1,7 +1,7 @@
 """Utils : Methods and Class which are used for tmessage."""
 from html.parser import HTMLParser
 from simple_chalk import chalk
-
+from cryptography.fernet import Fernet
 
 # pylint: disable=W0223
 class MessageParser(HTMLParser):
@@ -48,3 +48,28 @@ def get_formatted_message(raw_message):
     parser = MessageParser()
     parser.feed(raw_message)
     return parser.message_to_show
+
+
+def get_encryption_key():
+    """get the encryption key to encrypt and decrypt.
+
+        TODO:  need to avoid static key by implementing Diffie-Hellman.
+        # Fernet.generate_key() will generate a new key
+    """
+    return b'pxry9LY9AKkgHr-Zfs7x8hkUwJSHkdKxyeuq_retZ4A='
+
+
+def encrypt_message(message):
+    """encrypt a message and returns the encrypted message"""
+    key = get_encryption_key()
+    fernet = Fernet(key)
+    encrypted_message = fernet.encrypt(message)
+    return encrypted_message
+
+
+def decrypt_message(message):
+    """decrypt an encrypted message and returns the message"""
+    key = get_encryption_key()
+    fernet = Fernet(key)
+    decrypted_message = fernet.decrypt(message)
+    return decrypted_message