prof-oms/prof-oms.py

"""
Decrypt/Save or Encrypt/Upload attachments in OMEMO sessions. 
Build on "browser.py" skel (https://github.com/boothj5/profanity-plugins/blob/0.5.1/stable/browser.py)
"""

import prof
import requests
import magic
import re
import os
import tempfile

from os.path import expanduser
from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes

_links = {}
_lastlink = {}
iv = b''
key = b''
flag = False

def enc(key, iv, path_in):
    aes = AES.new(key, AES.MODE_GCM, nonce=iv)
    with open(path_in, "rb") as fin:
        with open(tempfile.gettempdir() + '/' + os.path.basename(path_in), 'wb') as fout:
            ftemp,tag = aes.encrypt_and_digest(fin.read())
            fout.write(ftemp + tag)

    strout = tempfile.gettempdir() + '/' + os.path.basename(path_in)
    return strout

def _get_url():
    link = None
    jid = prof.get_current_recipient()
    room = prof.get_current_muc()

    # check if in chat window
    if jid is not None:

        # check for link from recipient
        if jid in _lastlink.keys():
            link = _lastlink[jid]
        else:
            prof.cons_show("No links found from " + jid)

    # check if in muc window
    elif room is not None:
        if room in _lastlink.keys():
            link = _lastlink[room]
        else:
            prof.cons_show("No links found in " + room)

    # not in chat/muc window
    else:
        prof.cons_show("You must supply a aesgcm URI to the /oms command")

    return link

def _cmd_oms(arg1=None, arg2=None):
    global flag
    global iv
    global key
    global _lastlink

    link = None

    if arg1 == "setdir":
        if not arg2:
            prof.cons_bad_cmd_usage("/oms")
        else:
            prof.settings_string_set("oms", "download.path", arg2)
            prof.cons_show("oms.py set download path to" + arg2)

        return

    if arg1 == "send":
        if not arg2:
            prof.cons_bad_cmd_usage("/oms")
        else:
            key = get_random_bytes(32)
            iv = get_random_bytes(12)
            jid = prof.get_current_recipient()
            filein_enc = arg2
            strout = enc(key, iv, filein_enc)
            flag = True
            prof.send_line("/sendfile " + strout)

        return

    if arg1 == "get":
        if not arg2:
            prof.cons_bad_cmd_usage("/oms")
        else:
            link = arg2 if arg2 is not None else _get_url()

            if not link:
                prof.log_debug("oms.py: no aesgcm link found for oms command")
                return

            folder = prof.settings_string_get("oms", "download.path", "")
            if folder == "":
                prof.cons_show("Download path not set, see /help oms")
                return

            if folder[0] == '~':
                folder = expanduser("~") + folder[1:]

            address = re.sub("aesgcm","https",link.split('#')[0])
            if len(link.split('#')[1]) == 88:
                iv = bytes.fromhex(link.split('#')[1][:24])
                key = bytes.fromhex(link.split('#')[1][24:])
            elif len(link.split('#')[1]) == 96:
                iv = bytes.fromhex(link.split('#')[1][:32])
                key = bytes.fromhex(link.split('#')[1][32:])
            else:
                prof.chat_show(jid, "Unknown key length!")
                return
            

            response = requests.get(address)
            if response.status_code != 200:
                prof.log_debug("oms.py: received non 200 response")
                return
            
            file_c = response.content[:-16]
            tag = response.content[-16:]
            
            aes = AES.new(key, AES.MODE_GCM, nonce=iv)

            filename = address.split("/")[-1]
            full_path = folder + "/" + filename
            prof.log_debug("oms.py: Saving link {link} to file {file}".format(link=link, file=full_path))
            jid = prof.get_current_recipient()
            room = prof.get_current_muc()

            if jid is not None:
                prof.chat_show(jid, """ /\_/\\\n( o.o )  ♥ 1fd0 ♥\n > ^ <""")
                prof.chat_show(jid, "oms.py: Saving link to file {file}".format(file=full_path))
            elif room is not None:
                prof.room_show(room, """ /\_/\\\n( o.o )  ♥ 1fd0 ♥\n > ^ <""")
                prof.room_show(room, "oms.py: Saving link to file {file}".format(file=full_path))

            with open(full_path, 'wb') as f:
                f.write(aes.decrypt_and_verify(file_c, tag))

            return

def prof_init(version, status, account_name, fulljid):
    synopsis = [ 
        "/oms",
        "/oms get <url>",
        "/oms setdir <folder>",
        "/oms send <file_path>"
    ]
    description = (
        "/oms get: " +
        "Save a aesgcm link contents." + 
        "If no argument is supplied, the last aesgcm URI in the current chat or room will be used. " + 
        "Tab autocompletion will go through all previous links in the current chat/room. " +
        "/oms send: " +
        "Send encrypted file to chat recipient or muc."
    )
    args = [
        [ "setdir <dir>",     "Filesystem path to store downloaded content" ],
        [ "get <url>",        "URL from which download content" ],
        [ "send <file_path>", "Path of the file to upload" ]
    ]
    examples = [
        "/oms get aesgcm://accrocchio.org/Tsf5Qs98m/donatella_rettore.jpg#29deb53ff0328544cea03385175091cc402586992345dc3902658192bcba024921884728fd4f374922491fd0",
        "/oms setdir ~/Downloads",
        "/oms send /path/to/my/file",
    ]

    prof.register_command("/oms", 0, 2, synopsis, description, args, examples, _cmd_oms)
    prof.completer_add("/oms", [ "setdir", ])


def prof_on_chat_win_focus(jid):
    prof.completer_clear("/oms get")
    prof.completer_add("/oms", [ "setdir", "get", "send" ])
    prof.filepath_completer_add("/oms send")
    if jid in _links:
        prof.completer_add("/oms get", _links[jid])


def prof_on_room_win_focus(room):
    prof.completer_clear("/oms get")
    prof.completer_add("/oms", [ "setdir", "get", "send" ])
    prof.filepath_completer_add("/oms send")
    if room in _links:
        prof.completer_add("/oms get", _links[room])


def _process_message(jid, current_jid, message):
    links = re.findall(r'(aesgcm://\S+)', message)
    if len(links) > 0:
        if jid not in _links:
            _links[jid] = []
        
        # add to list of links for jid
        for link in links:
            if link not in _links[jid]:
                prof.log_debug("oms.py: Saving {link} for {jid}".format(link=link, jid=jid))
                _links[jid].append(link)
        
        # add to autocompleter if message for current window
        if current_jid == jid:
            prof.completer_add("/oms get", _links[jid])
        
        # set last link for jid
        _lastlink[jid] = links[len(links)-1]


def prof_post_chat_message_display(jid, resource, message):
    current_jid = prof.get_current_recipient()
    _process_message(jid, current_jid, message)


def prof_post_room_message_display(room, nick, message):
    current_jid = prof.get_current_muc()
    _process_message(room, current_jid, message)


def prof_on_room_history_message(room, nick, message, timestamp):
    current_jid = prof.get_current_muc()
    _process_message(room, current_jid, message)

def prof_pre_chat_message_send(jid, message):
    global flag
    if not flag:
        return message
    else:
        new_message = re.sub("^https", "aesgcm", message) + '#' + iv.hex() + key.hex()
        flag = False
        return new_message

def prof_pre_room_message_send(jid, message):
    global flag
    if not flag:
        return message
    else:
        new_message = re.sub("^https", "aesgcm", message) + '#' + iv.hex() + key.hex()
        flag = False
        return new_message