dru runmain

1 files changed, 56 insertions, 48 deletions

diff --git a/matrix-redact.py b/matrix-redact.py
index e7de4ea..578bb6d 100644
--- a/matrix-redact.py
+++ b/matrix-redact.py
@@ -11,6 +11,7 @@ import argparse
 import getpass
 import sys
 import subprocess
+from datetime import datetime
 from nio import(
         AsyncClient,
         MatrixRoom,
@@ -32,7 +33,7 @@ from nio import(
 
         responses
 )
-from typing import Optional
+from typing import Optional, List
 
 
 async def select_room(client: AsyncClient, sync_response) -> MatrixRoom:
@@ -65,19 +66,10 @@ async def client_login() -> Optional[AsyncClient]:
     print("login was successful, returning client.")
     return client
 
-# TODO: Can we fetch the total number of messages from `client.user_id` first, to show % progress?
-async def redact_room(client: AsyncClient, room: MatrixRoom) -> None:
-    print(f"\nAll events sent by '{client.user_id}' in " +
-          f"'{room['room_id']}' ({room['display_name']}) will be deleted.")
-    if not input("\n\nIs this room correct? (Type 'Continue') to proceed: ") == "Continue":
-        return None
-
+async def get_all_events(client: AsyncClient, room: MatrixRoom):
     sync_resp = await client.sync(timeout=30000, full_state=True)
-    start_token = sync_resp.rooms.join[room['room_id']].timeline.prev_batch
-    current_token = start_token
-
-    tracked = set()
-    events_redacted = 0
+    current_token = sync_resp.rooms.join[room['room_id']].timeline.prev_batch
+    events = []
 
     while True:
         try:
@@ -99,49 +91,63 @@ async def redact_room(client: AsyncClient, room: MatrixRoom) -> None:
             print("no more messages to fetch.")
             break
 
-        tasks = []
         for event in resp.chunk:
             event_id = event.event_id
 
-            if (event_id in tracked
-            or not hasattr(event, 'sender')
-            or event.sender != client.user_id
-            or isinstance(event, RedactionEvent)
-            or isinstance(event, RedactedEvent)
-            or isinstance(event, (RoomMemberEvent,
-                                 RoomCreateEvent,
-                                 RoomEncryptionEvent,
-                                 RoomGuestAccessEvent,
-                                 RoomHistoryVisibilityEvent,
-                                 RoomJoinRulesEvent,
-                                 PowerLevelsEvent))):
+            if (event_id in {e.event_id for e in events}
+                or not hasattr(event, 'sender')
+                or event.sender != client.user_id
+                or isinstance(event, RedactionEvent)
+                or isinstance(event, RedactedEvent)
+                or isinstance(event, (RoomMemberEvent,
+                                      RoomCreateEvent,
+                                      RoomEncryptionEvent,
+                                      RoomGuestAccessEvent,
+                                      RoomHistoryVisibilityEvent,
+                                      RoomJoinRulesEvent,
+                                      PowerLevelsEvent))):
                 continue
 
-            tracked.add(event_id)
+            events.append(event)
+    if len(events) > 1:
+        return events
+    return None
+
+async def redact_room(client: AsyncClient, room: MatrixRoom, dry_run) -> None:
+    print(f"\nAll events sent by '{client.user_id}' in " +
+          f"'{room['room_id']}' ({room['display_name']}) will be deleted.")
+    if not input("\n\nIs this room correct? (Type 'Continue') to proceed: ") == "Continue":
+        return None
+
+    events = await get_all_events(client, room)
+    events_redacted = 0
 
-            try:
-# NOTE:
-  # Fetching room events seems to not have a rate limit (or at least very minimally)
-  # Can we therefore in concurrency, fetch events "ahead" whenever redaction limits are returned?
+    if not events:
+        print("Could not get events.")
+        return None
+    print(f"Total number of events to redact: {len(events)}")
 
+    for event in events:
+        print(f"'{event.event_id}'")
+        try:
+            if not dry_run:
                 redact_resp = await client.room_redact(
                         room['room_id'],
-                        event_id,
+                        event.event_id,
                         reason=""
                 )
-
                 if isinstance(redact_resp, RoomRedactError):
-                    print(f"Could not redact event {event.event_id} in {room['room_id']}: {redact_resp}")
-                else:
-                    print(f"Successfully redacted event {event.event_id} in {room['room_id']}")
-                    events_redacted += 1
-                    print(f"current redactions: {events_redacted}")
-                    await asyncio.sleep(1)
-            except Exception as e:
-                print(f"Error redacting event {event.event_id}: {e}")
-
-    for event in tracked:
-        print(event)
+                    print(f"Could not redact event {event} in {room['room_id']}: {redact_resp}")
+                    continue
+
+            events_redacted += 1
+
+            print('\r', end='')
+            print(f"Progress: {events_redacted}/{len(events)}" +
+                  f" ({(events_redacted / len(events)) * 100:.2f}%)\t ",
+                  end='', flush=True)
+        except Exception as e:
+            print(f"Error redacting event {event}: {e}")
 
     if events_redacted:
         print(f"Total messages redacted: {events_redacted}")
@@ -169,7 +175,7 @@ async def main(args) -> None:
         room = await select_room(client, sync_resp)
 
         if room:
-            await redact_room(client, room)
+            await redact_room(client, room, args.dry)
 
     except Exception as e:
         print(f"Error: {e}")
@@ -184,8 +190,10 @@ if __name__ == "__main__":
             description="matrix-mcnt: Matrix Unread Message Count"
     )
 
-    # TODO:
-    # - delete just from a specific date range
-    # - delete the first a to b messages in DIRECTION
+    parser.add_argument(
+            "--dry",
+            help="Dry run (doesn't actually delete messages)",
+            action="store_true"
+    )
 
     asyncio.run(main(parser.parse_args()))