added ban, unban, and get commands

README.md

@@ -4,9 +4,10 @@ Manage instance blocks on mastodon
 
 Note: currently only import and export commands are supported. Also only tested with Glitch-Soc 3.0.1.
 
-PLEASE BACKUP YOUR DATABASE BEFORE USING TOOL
-
 ## Installation
+
+Minimum Python version: 3.6.0
+
 `pip install -r requirements.txt`
 
 `cp block.py path/to/mastodon/root`

block.py

@@ -1,4 +1,3 @@
- 
 #!/usr/bin/env python3
 '''
 Mastodon Block Manager v0.1
@@ -28,7 +27,7 @@ def dbconn():
 		print('Could not find ".env.production" in current directory.')
 		sys.exit(1)
 
-	dbhost = env.get('DB_HOST', 'localhost')
+	dbhost = env.get('DB_HOST', '/var/run/postgresql')
 	dbport = int(env.get('DB_PORT', 5432))
 	dbname = env.get('DB_NAME', 'mastodon_production')
 	dbuser = env.get('DB_USER', env.get('USER'))
@@ -57,9 +56,19 @@ def boolean(raw_val):
 		return val
 
 
-def dump(data):
-	db = dbconn()
+def ban_check(domain):
+	query = db.query(f'SELECT * FROM public.domain_blocks WHERE domain = \'{domain}\';').dictresult()
 
+	if query == []:
+		return
+
+	blocks = [domain['domain'] for domain in query]
+
+	if domain in blocks:
+		return True
+
+
+def dump(data):
 	if not data or not exists(dirname(data[0])):
 		print('Path doesn\'t exist. Saving to current directory instead as "block.csv"')
 		filename = 'block.csv'
@@ -84,7 +93,6 @@ def dump(data):
 				domain['public_comment']
 			])
 
-	db.close()
 	return 'Done! :3'
 
 
@@ -94,8 +102,6 @@ def load(data):
 	1 suspend
 	2 none
 	'''
-	db = dbconn()
-
 	if not data or not exists(data[0]):
 		print('Path doesn\'t exist. Saving to current directory instead as "block.csv"')
 		filename = 'block.csv'
@@ -148,22 +154,86 @@ def load(data):
 			)
 			print(f'Created new block for {row[0]}')
 
-	db.close()
 	return 'Done! :3'
 
 
+def ban(data):
+	if not data:
+		return 'Missing domain to block'
+
+	if not ban_check(data[0]):
+		date = datetime.datetime.now()
+
+		if data[1] in ['suspend', 'silence', 'none']:
+			bantype = {'silence': 0, 'suspend': 1, 'none': 2}
+			severity = bantype[data[1]]
+
+		else:
+			severity = 0
+
+		db.insert('public.domain_blocks',
+			domain=data[0],
+			severity=severity,
+			reject_media=data[2] if data[2] != None else False,
+			reject_reports=data[3] if data[3] != None else False,
+			private_comment=data[4] if boolean(data[4]) else None,
+			public_comment=data[5] if boolean(data[5]) else None,
+			created_at=date,
+			updated_at=date
+		)
+		return f'Created new block for {data[0]}'
+
+	# This will eventually be replaced with code to update the ban
+	return f'{data[0]} was already banned'
+
+
+def unban(data):
+	if not data:
+		return 'Missing domain to block'
+
+	if not ban_check(data[0]):
+		return f'{data[0]} wasn\'t in the banlist'
+
+	rowquery = db.query(f'SELECT * FROM public.domain_blocks WHERE domain = \'{data[0]}\';').dictresult()
+	rowdata = rowquery[0] if rowquery else None
+
+	if db.delete('public.domain_blocks', {'id': rowdata['id']}) == 1:
+		return f'Unbanned {data[0]}'
+
+	else:
+		return f'Failed to unban {data[0]}'
+
+
+def get(data):
+	domains = db.query('SELECT * FROM public.domain_blocks;').dictresult()
+
+	msg = 'Domain	Severity	reject media	reject reports	private comment	public comment'
+
+	for domain in domains:
+		msg += f"\n{domain['domain']}\t{domain['severity']}\t{boolean(domain['reject_media'])}\t{boolean(domain['reject_reports'])}\t{domain['private_comment']}\t{domain['public_comment']}"
+
+	return msg
+
+
 def main():
 	arg = sys.argv
 
 	cmdhelp = '''Mastodon Blocklist Manager
 
-import <file>	Import a blocklist
-export <file>	Export the blocklist to csv
+import [file]
+export [file]
+ban <domain> [severity] [reject media] [reject reports] [private commend] [public comment]
+unban <domain>
+update <domain> [severity] [reject media] [reject reports] [private commend] [public comment]
+get
 '''
 
 	tasks = {
 		'export': dump,
-		'import': load
+		'import': load,
+		'ban': ban,
+		'unban': unban,
+		'get': get
 	}
 
 	if len(arg) < 2 or arg[1] not in tasks:
@@ -175,4 +245,6 @@ export <file>	Export the blocklist to csv
 
 
 if __name__ == '__main__':
+	db = dbconn()
 	print(main())
+	db.close()