Added menu to create, delete or list bots

bots-setup.py

@@ -1,12 +1,27 @@
 import getpass
 import os
+import time
 import sys
 import psycopg2
 from mastodon import Mastodon
 from mastodon.Mastodon import MastodonMalformedEventError, MastodonNetworkError, MastodonReadTimeout, MastodonAPIError, MastodonIllegalArgumentError
 import uuid
+import requests
+from prettytable import PrettyTable
 import pdb
 
+menu_options = {
+    1: 'New Bot',
+    2: 'Delete Bot',
+    3: 'List Bots',
+    4: 'Exit',
+}
+
+def print_menu():
+
+    for key in menu_options.keys():
+        print (key, '-', menu_options[key] )
+
 def get_config():
 
     # Load configuration from config file
@@ -30,18 +45,36 @@ def check_db_conn():
 
         sys.exit("Exiting. Run 'python db-setup' again")
 
+def check_host(url):
+
+    host_exists = False
+
+    try:
+
+        r = requests.get(f'https://{url}')
+
+        host_exists = True
+
+    except requests.exceptions.RequestException as error:
+
+        print(f'\n{url}: {error}')
+
+        time.sleep(3)
+
+    return host_exists
+
 class Bot:
 
     name = "Bot"
 
-    def __init__(self, hostname, software, username, password, app_name, twitter_username):
+    def __init__(self):
 
-        self.hostname = hostname
+        self.hostname = input("\nEnter Mastodon/Pleroma hostname: ")
-        self.software = software
+        self.software = input("Server software (mastodon or pleroma)? ")
-        self.username = username
+        self.username = input(f'User name, ex. user@{self.hostname}? ')
-        self.password = password
+        self.password = getpass.getpass("User password? ")
-        self.app_name = app_name
+        self.app_name = 'replicator'
-        self.twitter_username = twitter_username
+        self.twitter_username = input(f'Twitter username (ex. jack)? ')
 
     def log_in(self):
 
@@ -131,9 +164,9 @@ class Bot:
 
         save_error = None
 
-        mastodon_sql = f'INSERT INTO maccounts (bot_id, hostname, username, client_id, client_secret, client_token, hostname_soft) VALUES (%s,%s,%s,%s,%s,%s,%s)'
+        mastodon_sql = 'INSERT INTO maccounts (bot_id, hostname, username, client_id, client_secret, client_token, hostname_soft) VALUES (%s,%s,%s,%s,%s,%s,%s)'
 
-        twitter_sql = f'INSERT INTO taccounts (bot_id, username) VALUES (%s,%s)'
+        twitter_sql = 'INSERT INTO taccounts (bot_id, username) VALUES (%s,%s)'
 
         conn = None
 
@@ -163,6 +196,146 @@ class Bot:
 
         return save_error
 
+class GetBot:
+
+    name = "GetBot"
+
+    def __init__(self):
+
+        self.twitter_username = input(f'Twitter username? ')
+
+    def delete_account(self):
+
+        del_error = None
+
+        twitter_sql = 'SELECT bot_id from taccounts where username=(%s)'
+
+        twitter_del_sql = 'DELETE FROM taccounts WHERE username=(%s)'
+
+        mastodon_del_sql = 'DELETE FROM maccounts WHERE bot_id=(%s)'
+
+        conn = None
+
+        try:
+
+            conn = psycopg2.connect(database = replicator_db, user = replicator_db_user, password = "", host = "/var/run/postgresql", port = "5432")
+
+            cur = conn.cursor()
+
+            cur.execute(twitter_sql, (self.twitter_username,))
+
+            row = cur.fetchone()
+
+            if row != None:
+
+                delete_it = confirm_deletion(self.twitter_username)
+
+                if delete_it:
+
+                    bot_id = row[0]
+
+                    cur.execute(twitter_del_sql, (self.twitter_username,))
+
+                    cur.execute(mastodon_del_sql, (bot_id,))
+
+                    conn.commit()
+
+                    print(f'{self.twitter_username} had been deleted.')
+
+                    time.sleep(2)
+
+            else:
+
+                print(f'{self.twitter_username} does not exists!')
+
+                time.sleep(2)
+
+        except (Exception, psycopg2.DatabaseError) as error:
+
+            del_error = error.pgcode
+
+            sys.stdout.write(f'\n{str(error)}\n')
+
+        finally:
+
+            if conn is not None:
+
+                conn.close()
+
+        return (del_error)
+
+    def print(self):
+
+        print(f'\nBot found: {self.twitter_username}\n')
+
+def confirm_deletion(twitter_username):
+
+    answer = ''
+
+    while (answer != 'Yes' and answer != 'No'):
+
+        answer = input(f'Are you sure deleting {twitter_username}?: ')
+
+        if answer == 'Yes':
+
+            print('Ok, deleting it...')
+
+            delete_it = True
+
+        if answer == 'No':
+
+            print('Doing nothing.')
+
+            delete_it = False
+
+    return delete_it
+
+def list_bots():
+
+    bots_list = []
+
+    t_usernames_lst = []
+
+    mastodon_sql = 'SELECT bot_id, hostname, username, hostname_soft FROM maccounts'
+
+    twitter_sql = 'SELECT username from taccounts where bot_id=(%s)'
+
+    conn = None
+
+    try:
+
+        conn = psycopg2.connect(database = replicator_db, user = replicator_db_user, password = "", host = "/var/run/postgresql", port = "5432")
+
+        cur = conn.cursor()
+
+        cur.execute(mastodon_sql)
+
+        rows = cur.fetchall()
+
+        for row in rows:
+
+            cur.execute(twitter_sql, (row[0],))
+
+            username_row = cur.fetchone()
+
+            t_usernames_lst.append(username_row[0])
+
+            bots_list.append(row)
+
+    except (Exception, psycopg2.DatabaseError) as error:
+
+        get_error = error.pgcode
+
+        sys.stdout.write(f'\n{str(error)}\n')
+
+    finally:
+
+        if conn is not None:
+
+            conn.close()
+
+    return (bots_list, t_usernames_lst)
+
 def get_db_params( parameter, file_path ):
 
     if not os.path.isfile(file_path):
@@ -185,41 +358,99 @@ def get_db_params( parameter, file_path ):
 
 if __name__ == '__main__':
 
-    config_filepath, replicator_db, replicator_db_user = get_config()
+    while(True):
 
-    check_db_conn()
+        print('\n')
 
-    bot_hostname = input("Enter Mastodon/Pleroma hostname: ")
+        print_menu()
-    bot_software = input("Server software (mastodon or pleroma)? ")
-    bot_username = input(f'User name, ex. user@{bot_hostname}? ')
-    bot_password = getpass.getpass("User password? ")
-    bot_app_name = 'replicator'
-    bot_twitter_username = input(f'Twitter username (ex. jack)? ')
-    
-    myBot = Bot(bot_hostname, bot_software, bot_username, bot_password, bot_app_name, bot_twitter_username)
 
-    is_duplicate = myBot.check_account()
+        option = ''
 
-    if is_duplicate:
+        try:
 
-        print(f'Bot account already exist!')
+            option = int(input('\nEnter your choice: '))
 
-        sys.exit()
+        except:
 
-    else:
+            print('\nWrong input. Please enter a number between 1 and 4.\n')
 
-        client_id, client_secret, token = myBot.log_in()
+        if option == 1:
 
-        if len(token) > 0:
+            config_filepath, replicator_db, replicator_db_user = get_config()
 
-            save_error = myBot.save_account(client_id, client_secret, token)
+            check_db_conn()
 
-            if save_error == None:
+            myBot = Bot()
 
-                print(f'Done.')
+            host_exists = check_host(myBot.hostname)
 
-            else:
+            if host_exists:
 
-                if save_error == '423505':
+                is_duplicate = myBot.check_account()
+
+                if is_duplicate:
+
+                    print(f'\nBot account already exist!\n')
+
+                    time.sleep(3)
+
+                else:
+
+                    client_id, client_secret, token = myBot.log_in()
+
+                    if len(token) > 0:
+
+                        save_error = myBot.save_account(client_id, client_secret, token)
+
+                        if save_error == None:
+
+                            print(f'Bot added!')
+
+                            time.sleep(3)
+
+                        else:
+
+                            if save_error == '423505':
+
+                                print(f'error: {save_error}')
+
+        elif option == 2:
+
+            config_filepath, replicator_db, replicator_db_user = get_config()
+
+            check_db_conn()
+
+            getbot = GetBot()
+
+            getbot.delete_account()
+
+        elif option == 3:
+
+            config_filepath, replicator_db, replicator_db_user = get_config()
+
+            check_db_conn()
+
+            bots_list, t_usernames_lst = list_bots()
+
+            print_table = PrettyTable()
+
+
+            print_table.field_names = ['No.','Twitter username', 'bot account', 'host', 'host software']
+
+            i = 0
+
+            for bot in bots_list:
+
+                print_table.add_row([str(i+1), t_usernames_lst[i], bot[2], bot[1], bot[3]])
+
+                i += 1
+
+            print(print_table)
+
+            print(f'Total: {len(bots_list)}')
+
+        elif option == 4:
+
+            print('Bye!')
+            exit()
 
-                    sys.exit()

requirements.txt

@@ -1,3 +1,4 @@
 psycopg2>=2.9.3
 Mastodon.py>=1.5.1
 tweepy>=4.5.0
+PrettyTable>=3.1.1