handle suspended/silenced accounts

mastoblocks.py

@@ -1,13 +1,16 @@
 #!/usr/bin/env python3
 '''
-Mastodon Block Manager v0.1
+Mastodon Block Manager v0.2
 by Izalia Mae (@Izalia@barkshark.xyz)
 
-requirements:
+installation:
+	sudo apt install python3-dev libpq-dev
 	pip3 install pygresql envbash
+	Copy to mastodon directory
 '''
-import datetime, json, os, sys
+import json, os, sys
 
+from datetime import datetime
 from os import environ as env
 from pathlib import Path
 
@@ -21,10 +24,11 @@ severity_import = {
 	'none': 2
 }
 
-severity_export = {
-	0: 'silence',
-	1: 'suspend',
-	2: 'none'
+severity_export = {v:k for k,v in severity_import.items()}
+
+block_type = {
+	'suspend': 'suspended_at',
+	'silence': 'silenced_at'
 }
 
 
@@ -51,8 +55,6 @@ def dbconn():
 
 
 def dump(data):
-	db = dbconn()
-
 	if not data:
 		print('Path doesn\'t exist. Saving to current directory instead as "block.json"')
 		filename = Path('blocks.json').resolve()
@@ -75,13 +77,10 @@ def dump(data):
 			new_data.append(domain)
 		fd.write(json.dumps(new_data, indent=4))
 
-	db.close()
 	return 'Done! :3'
 
 
 def load(data):
-	db = dbconn()
-
 	if not data:
 		print('Path doesn\'t exist. Reading from "blocks.json" instead.')
 		filename = Path('blocks.json').resolve()
@@ -95,7 +94,7 @@ def load(data):
 		domain = line['domain']
 		rowquery = db.query(f"SELECT * FROM public.domain_blocks WHERE domain = '{domain}';").dictresult()
 		row = rowquery[0] if rowquery else None
-		date = datetime.datetime.now()
+		date = datetime.now()
 
 		data = {
 			'severity': severity_import.get(line['severity'], 1),
@@ -115,17 +114,34 @@ def load(data):
 			db.insert('public.domain_blocks', **data)
 			print(f'Created new block for {line["domain"]}')
 
-	db.close()
+		btype = block_type.get(line['severity'])
+
+		if btype:
+			accts = db.query(f"SELECT * FROM public.accounts WHERE domain LIKE '%{domain}'").dictresult()
+
+			for acct in accts:
+				account_id = acct['id']
+				db.update('public.accounts', {'id': account_id}, **{btype: date})
+
+				if btype == 'suspended_at':
+					follows = db.query(f"SELECT * FROM public.follows WHERE account_id = '{account_id}' or target_account_id = '{account_id}'")
+
+					for follow in follows:
+						db.delete('public.follows', follow)
+
 	return 'Done! :3'
 
 
+db = dbconn()
+
+
 def main():
 	arg = sys.argv
 
 	cmdhelp = '''Mastodon Blocklist Manager
 
-import <file>	Import a blocklist
-export <file>	Export the blocklist to csv
+import <file>	Import a json blocklist
+export <file>	Export the blocklist to json
 '''
 
 	tasks = {
@@ -136,8 +152,17 @@ export <file>	Export the blocklist to csv
 	if len(arg) < 2 or arg[1] not in tasks:
 		return cmdhelp
 
-	msg = tasks[arg[1]](None if len(arg) < 3 else arg[2:])
+	try:
+		db.begin()
+		msg = tasks[arg[1]](None if len(arg) < 3 else arg[2:])
 
+	except Exception as e:
+		db.rollback()
+		db.close()
+		raise e from None
+
+	db.commit()
+	db.close()
 	return msg