Refactored to use Bot Class

bots-setup.py

@@ -7,80 +7,6 @@ from mastodon.Mastodon import MastodonMalformedEventError, MastodonNetworkError,
 import uuid
 import pdb
 
-def check_account(username):
-
-    is_duplicate = False
-
-    sql = 'select username from maccounts where username=(%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(sql, (username,))
-
-        row = cur.fetchone()
-
-        if row != None:
-
-            is_duplicate = True
-
-    except (Exception, psycopg2.DatabaseError) as error:
-
-        check_error = error.pgcode
-
-        sys.stdout.write(f'\n{str(error)}\n')
-
-    finally:
-
-        if conn is not None:
-
-            conn.close()
-
-        return is_duplicate
-
-def save_bot_account(hostname, username, client_id, client_secret, token, twitter_username):
-
-    bot_uuid = str(uuid.uuid4())
-
-    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)'
-
-    twitter_sql = f'INSERT INTO taccounts (bot_id, username) VALUES (%s,%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, (bot_uuid, hostname, username, client_id, client_secret, token, hostname_soft))
-
-        cur.execute(twitter_sql, (bot_uuid, twitter_username))
-
-        conn.commit()
-
-    except (Exception, psycopg2.DatabaseError) as error:
-
-        save_error = error.pgcode
-
-        sys.stdout.write(f'\n{str(error)}\n')
-
-    finally:
-
-        if conn is not None:
-
-            conn.close()
-
-    return save_error
-
 def get_config():
 
     # Load configuration from config file
@@ -104,42 +30,45 @@ def check_db_conn():
 
         sys.exit("Exiting. Run 'python db-setup' again")
 
-def ask_account():
+class Bot:
 
-    print(f'Setting up Mastodon/Pleroma bot account...\n')
+    name = "Bot"
-    hostname = input("Enter Mastodon/Pleroma hostname: ")
-    hostname_soft = input("Server software (mastodon or pleroma)? ")
-    user_name = input(f'User name, ex. user@{hostname}? ')
-    user_password = getpass.getpass("User password? ")
-    app_name = 'replicator'
-    twitter_username = input(f'Twitter username (ex. jack)? ')
 
-    return (hostname, hostname_soft, user_name, user_password, app_name, twitter_username)
+    def __init__(self, hostname, software, username, password, app_name, twitter_username):
 
-def log_in(hostname, user_name, user_password, app_name):
+        self.hostname = hostname
+        self.software = software
+        self.username = username
+        self.password = password
+        self.app_name = app_name
+        self.twitter_username = twitter_username
+
+    def log_in(self):
 
         token = ''
 
         try:
 
             response = Mastodon.create_app(
-            app_name,
+                self.app_name,
                 scopes=["read","write"],
                 to_file=None,
-            api_base_url=hostname
+                api_base_url=self.hostname
                 )
             client_id = response[0]
             client_secret = response[1]
 
-        mastodon = Mastodon(client_id = client_id, client_secret = client_secret, api_base_url = hostname)
+            mastodon = Mastodon(client_id = client_id, client_secret = client_secret, api_base_url = self.hostname)
         
             token = mastodon.log_in(
-            user_name,
+                self.username,
-            user_password,
+                self.password,
                 scopes = ["read", "write"],
                 to_file = None
             )
 
+            print('Log in succesful!')
+
         except MastodonIllegalArgumentError as i_error:
 
             sys.stdout.write(f'\n{str(i_error)}\n')
@@ -160,6 +89,80 @@ def log_in(hostname, user_name, user_password, app_name):
 
             return (client_id, client_secret, token)
 
+    def check_account(self):
+
+        is_duplicate = False
+
+        sql = 'select username from maccounts where username=(%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(sql, (self.username,))
+
+            row = cur.fetchone()
+
+            if row != None:
+
+                is_duplicate = True
+
+        except (Exception, psycopg2.DatabaseError) as error:
+
+            check_error = error.pgcode
+
+            sys.stdout.write(f'\n{str(error)}\n')
+
+        finally:
+
+            if conn is not None:
+
+                conn.close()
+
+            return is_duplicate
+
+    def save_account(self, client_id, client_secret, token):
+
+        bot_uuid = str(uuid.uuid4())
+
+        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)'
+
+        twitter_sql = f'INSERT INTO taccounts (bot_id, username) VALUES (%s,%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, (bot_uuid, self.hostname, self.username, client_id, client_secret, token, self.software))
+
+            cur.execute(twitter_sql, (bot_uuid, self.twitter_username))
+
+            conn.commit()
+
+        except (Exception, psycopg2.DatabaseError) as error:
+
+            save_error = error.pgcode
+
+            sys.stdout.write(f'\n{str(error)}\n')
+
+        finally:
+
+            if conn is not None:
+
+                conn.close()
+
+        return save_error
+
 def get_db_params( parameter, file_path ):
 
     if not os.path.isfile(file_path):
@@ -186,9 +189,16 @@ if __name__ == '__main__':
 
     check_db_conn()
 
-    hostname, hostname_soft, user_name, user_password, app_name, twitter_username = ask_account()
+    bot_hostname = input("Enter Mastodon/Pleroma hostname: ")
+    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)? ')
     
-    is_duplicate = check_account(user_name)
+    myBot = Bot(bot_hostname, bot_software, bot_username, bot_password, bot_app_name, bot_twitter_username)
+
+    is_duplicate = myBot.check_account()
 
     if is_duplicate:
 
@@ -198,11 +208,11 @@ if __name__ == '__main__':
 
     else:
 
-        client_id, client_secret, token = log_in(hostname, user_name, user_password, app_name)
+        client_id, client_secret, token = myBot.log_in()
 
         if len(token) > 0:
 
-            save_error = save_bot_account(hostname, hostname_soft, user_name, client_id, client_secret, token, twitter_username)
+            save_error = myBot.save_account(client_id, client_secret, token)
 
             if save_error == None: