pubnub_e2e_chat

pubnub_e2e_chat.py

# "I don't trust e2e chats, I don't know what's under the hood on signal, actually"
#
# "Ok, gimme 10 minutes, let me setup a e2e chat script. We'll use a free service on the 
# other side of the world, transferring meaningless base64 AES encrypted payloads over the net"
#
# "Ah, cool".
#
# Very buggy - Don't expect fixes
#
# Signup for api keys at https://pubnub.com
#
# pip install PyCrypto ; pip install pubnub
#
import json
import sys
import platform

print('Loading')

from pubnub.callbacks import SubscribeCallback
from pubnub.pnconfiguration import PNConfiguration
from pubnub.pubnub import PubNub
from pprint import pprint
import base64
import hashlib

import os
from Crypto.Cipher import AES


def get_keys():
    try:
        import os
        __location__ = os.path.realpath(
            os.path.join(os.getcwd(), os.path.dirname(__file__)))
        with open('{}/.chatkeys'.format(__location__), 'r') as f:
            keys = json.load(f)
            return keys
    except:
        return None


def save_keys(pub, sec):
    try:
        import os
        __location__ = os.path.realpath(
            os.path.join(os.getcwd(), os.path.dirname(__file__)))
        with open('{}/.chatkeys'.format(__location__), 'w') as f:
            json.dump([pub, sec], f)
            return True
    except:
        return None




class AESCipherService(object):

    def __init__(self, key, iv=None):
        self.bs = 32
        self.key = hashlib.sha256(key).digest()
        self.iv = iv

    def encrypt(self, raw):
        raw = self._pad(raw)
        iv = self.iv or os.urandom(16)
        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        if 'inux' in platform.platform():
            return base64.b64encode(iv + cipher.encrypt(raw)).decode()
        else:
            return base64.b64encode(iv + cipher.encrypt(raw.encode())).decode()

    def decrypt(self, enc):
        enc = base64.b64decode(enc)
        iv = enc[:AES.block_size]
        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        return self._unpad(cipher.decrypt(enc[AES.block_size:])).decode('utf-8')

    def _pad(self, s):
        return s + (self.bs - len(s) % self.bs) * chr(self.bs - len(s) % self.bs)

    @staticmethod
    def _unpad(s):
        return s[:-ord(s[len(s)-1:])]


class MySubscribeCallback(SubscribeCallback):
    def status(self, pubnub, status):
        pass

    def presence(self, pubnub, presence):
        pprint(presence.__dict__)

    def message(self, pubnub, message):
        message = message.__dict__

        try:
            rawmsg = encryption.decrypt(message['message']['content'])
            assert rawmsg
            rnickname, msg = rawmsg.split('|{}|'.format(separator))
            if rnickname != nickname:
                print('<{}> {}'.format(rnickname, msg))
        except:
            print('- UNKNOWN MESSAGE - CHECK SEED -')


pnconfig = PNConfiguration()

if __name__ == '__main__':
    keys = get_keys()
    kc = False
    if not keys:
        pnconfig.publish_key = input('Enter pubkey> ').strip()
        pnconfig.subscribe_key = input('Enter subkey> ').strip()
        kc = True
    else:
        print('Using pubkey {}\n'.format(keys[0]))
        ok = input('Y(default)/N >').strip().upper()
        ok = ok or 'Y'
        if ok != 'Y':
            pnconfig.publish_key = input('Enter pubkey> ').strip()
            pnconfig.subscribe_key = input('Enter subkey> ').strip()
            kc = True
        else:
            pnconfig.publish_key = keys[0]
            pnconfig.subscribe_key = keys[1]
    if kc:
        save_keys(pnconfig.publish_key, pnconfig.subscribe_key)
    pubnub = PubNub(pnconfig)
    pubnub.add_listener(MySubscribeCallback())
    nickname = None
    while not nickname:
        nickname = input('Enter your nickname> ').strip()
    seed = None
    while not seed:
        seed = input('Enter your seed> ').strip()
    encryption = AESCipherService(hashlib.sha256(seed.encode()).digest())
    separator = seed
    for x in range(0, 16):
        separator = hashlib.sha256(seed.encode()).hexdigest()
    pubnub.subscribe().channels("pubnub_onboarding_channel").with_presence().execute()
    while 1:
        try:
            message = input('<{}> '.format(nickname))
            if message == '/quit':
                print('Quit!')
                quit()
                sys.exit(-1)
            if message:
                try:
                    pubnub.publish().channel("pubnub_onboarding_channel").message(
                        {"sender": pnconfig.uuid, "content": encryption.encrypt(
                            '{}|{}|{}'.format(
                                nickname,
                                separator,
                                message
                            )
                        )
                         }
                    ).pn_async(lambda *a: True)
                except Exception as e:
                    print('Error: %s - Maybe try again' % e)
        except KeyboardInterrupt:
            print('Exiting')
            quit()