weechat-matrix/contrib/matrix_upload

#!/usr/bin/python3 -u
# Copyright © 2018 Damir Jelić <poljar@termina.org.uk>
#
# Permission to use, copy, modify, and/or distribute this software for
# any purpose with or without fee is hereby granted, provided that the
# above copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
# SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER
# RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF
# CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
# CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.


import os
import json
import magic
import requests
import argparse
from urllib.parse import urlparse
import urllib3

from nio import Api, UploadResponse, UploadError
from json.decoder import JSONDecodeError

urllib3.disable_warnings()

mime = magic.Magic(mime=True)


class Upload(object):
    def __init__(self, file, chunksize=1 << 13):
        self.file = file
        self.filename = os.path.basename(file)
        self.chunksize = chunksize
        self.totalsize = os.path.getsize(file)
        self.mimetype = mime.from_file(file)
        self.readsofar = 0

    def send_progress(self):
        message = {
            "type": "progress",
            "data": self.readsofar
        }
        to_stdout(message)

    def __iter__(self):
        with open(self.filename, 'rb') as file:
            while True:
                data = file.read(self.chunksize)

                if not data:
                    break

                self.readsofar += len(data)
                self.send_progress()

                yield data

    def __len__(self):
        return self.totalsize


class IterableToFileAdapter(object):
    def __init__(self, iterable):
        self.iterator = iter(iterable)
        self.length = len(iterable)

    def read(self, size=-1):
        return next(self.iterator, b'')

    def __len__(self):
        return self.length


def to_stdout(message):
    print(json.dumps(message), flush=True)


def error(e):
    message = {
        "type": "status",
        "status": "error",
        "message": str(e)
    }
    to_stdout(message)
    os.sys.exit()


def upload_process(args):
    file_path = os.path.expanduser(args.file)

    try:
        upload = Upload(file_path, 10)
    except (FileNotFoundError, OSError, IOError) as e:
        error(e)

    try:
        url = urlparse(args.homeserver)
    except ValueError as e:
        error(e)

    upload_url = ("https://{}".format(args.homeserver)
                  if not url.scheme else args.homeserver)
    _, api_path, _ = Api.upload(args.access_token, upload.filename)
    upload_url += api_path

    headers = {
        "Content-type": upload.mimetype,
    }

    proxies = {}

    if args.proxy_address:
        user = args.proxy_user or ""

        if args.proxy_password:
            user += ":{}".format(args.proxy_password)

        if user:
            user += "@"

        proxies = {
            "https": "{}://{}{}:{}/".format(
                args.proxy_type,
                user,
                args.proxy_address,
                args.proxy_port
            )
        }

    message = {
        "type": "status",
        "status": "started",
        "total": upload.totalsize,
        "mimetype": upload.mimetype,
        "file_name": upload.filename,
    }
    to_stdout(message)

    session = requests.Session()
    session.trust_env = False

    try:
        r = session.post(
            url=upload_url,
            auth=None,
            headers=headers,
            data=IterableToFileAdapter(upload),
            verify=(not args.insecure),
            proxies=proxies
        )
    except (requests.exceptions.RequestException, OSError) as e:
        error(e)

    try:
        json_response = json.loads(r.content)
    except JSONDecodeError:
        error(r.content)

    response = UploadResponse.from_dict(json_response)

    if isinstance(response, UploadError):
        error(str(response))

    message = {
        "type": "status",
        "status": "done",
        "url": response.content_uri
    }

    to_stdout(message)

    return 0


def main():
    parser = argparse.ArgumentParser(
        description="Download and decrypt matrix attachments"
    )
    parser.add_argument("file", help="the file that will be uploaded")
    parser.add_argument(
        "homeserver",
        type=str,
        help="the address of the homeserver"
    )
    parser.add_argument(
        "access_token",
        type=str,
        help="the access token to use for the upload"
    )
    parser.add_argument(
        "--encrypt",
        action="store_const",
        const=True,
        default=False,
        help="encrypt the file before uploading it"
    )
    parser.add_argument(
        "--insecure",
        action="store_const",
        const=True,
        default=False,
        help="disable SSL certificate verification"
    )
    parser.add_argument(
        "--proxy-type",
        choices=[
            "http",
            "socks4",
            "socks5"
        ],
        default="http",
        help="type of the proxy that will be used to establish a connection"
    )
    parser.add_argument(
        "--proxy-address",
        type=str,
        help="address of the proxy that will be used to establish a connection"
    )
    parser.add_argument(
        "--proxy-port",
        type=int,
        default=8080,
        help="port of the proxy that will be used to establish a connection"
    )
    parser.add_argument(
        "--proxy-user",
        type=str,
        help="user that will be used for authentication on the proxy"
    )
    parser.add_argument(
        "--proxy-password",
        type=str,
        help="password that will be used for authentication on the proxy"
    )

    args = parser.parse_args()
    upload_process(args)


if __name__ == "__main__":
    main()