#!/usr/bin/env python3
import argparse
import requests
import sys
from pprint import pformat
from time import sleep, time


# configure anti-throtle
sleep_for_seconds = 3
pagination_limit = 150


# parse arguments
parser = argparse.ArgumentParser()
parser.add_argument(
            "room_ids",
            help="room ID's",
            nargs='*')
parser.add_argument(
            "-n",
            "--nuke",
            help="try and delete all messages in the room, not only mine",
            action='store_true')
req_args = parser.add_argument_group('required')
req_args.add_argument(
            "-s",
            "--homeserver",
            help="homeserver's domain name",
            required=True)
req_args.add_argument(
            "-t",
            "--access-token",
            help="Matrix access token",
            required=True)
args = parser.parse_args()
homeserver = args.homeserver
access_token = args.access_token
room_ids = args.room_ids


def whoami(acces_token):
    headers = { "Authorization": f"Bearer {access_token}" }
    url = f"https://{homeserver}/_matrix/client/v3/account/whoami"
    response = requests.get(url, headers=headers)
    return response.json()['user_id']


def get_all_encrypted_messages(homeserver, room_id, access_token):
    my_username = whoami(access_token)
    print("My user:", my_username)
    messages = []
    next_batch = None
    while True:
        print("Reading messages...")
        # don't get throttled
        sleep(sleep_for_seconds)
        # define request parameters
        filter_param = '"types":["m.room.encrypted", "m.room.message"]'
        if not args.nuke:
            filter_param = '"senders":["' + my_username + '"],' + filter_param
        params = {
            "dir": "b",
            "limit": pagination_limit,
            "filter": '{' + filter_param + '}'
        }
        if next_batch:
            params["from"] = next_batch
        # make request
        url = f"https://{homeserver}/_matrix/client/v3/rooms/{room_id}/messages"
        headers = { "Authorization": f"Bearer {access_token}" }
        response = requests.get(url, params=params, headers=headers)
        response.raise_for_status()
        chunk = response.json()["chunk"]
        messages.extend(chunk)
        next_batch = response.json().get("end")
        if not next_batch:
            break
    return messages


def redact_message(message):
    transaction_id = 'm' + str(time()).replace('.', '')[:14] + '.0'
    event_id = message['event_id']
    h = homeserver
    r = room_id
    e = event_id
    t = transaction_id
    url = f"https://{h}/_matrix/client/r0/rooms/{r}/redact/{e}/{t}"
    headers = {
        "Authorization": f"Bearer {access_token}"
    }
    # don't get throttled
    sleep(sleep_for_seconds)
    response = requests.put(url, json={"reason":""}, headers=headers)
    print(pformat(response))
    response.raise_for_status()


for room_id in room_ids:
    all_messages = get_all_encrypted_messages(homeserver, room_id, access_token)
    message_count = 0
    for message in all_messages:
        if message['content']:
            redact_message(message)
            message_count += 1
            print("Redacted messages:", message_count)