spamcheck/spamcheck.py

from datetime import date, datetime, timedelta
#from mastodon import Mastodon
import time
import os
import json
import sys
import os.path
import operator
import psycopg2
from psycopg2 import sql
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT

class Spamcheck:

    name = "Spamcheck for Mastodon social servers"

    def __init__(self, mastodon_hostname=None, mastodon_db=None, mastodon_db_user=None, spamcheck_db=None, spamcheck_db_user=None):

        self.config_file = 'config/config.txt'

        is_setup = self.__check_setup(self)

        if is_setup:

            self.mastodon_hostname = self.__get_parameter("mastodon_hostname", self.config_file)
            self.mastodon_db = self.__get_parameter("mastodon_db", self.config_file)
            self.mastodon_db_user = self.__get_parameter("mastodon_db_user", self.config_file)
            self.spamcheck_db = self.__get_parameter("spamcheck_db", self.config_file)
            self.spamcheck_db_user = self.__get_parameter("spamcheck_db_user", self.config_file)
        
        else:

            self.mastodon_hostname, self.mastodon_db, self.mastodon_db_user, self.spamcheck_db, self.spamcheck_db_user = self.__setup(self)

        db_setup = self.__check_dbsetup(self)

        if not db_setup:

            self.__createdb(self)

    def new_registers(self, created_at_lst=[], id_lst=[], email_lst=[], ip_lst=[]):
    
        try:

            conn = None

            conn = psycopg2.connect(database = self.mastodon_db, user = self.mastodon_db_user, password = "", host = "/var/run/postgresql", port = "5432")

            cur = conn.cursor()

            cur.execute("select users.created_at, users.id, users.email, users.sign_up_ip from users where users.approved and users.created_at > now() - interval '7 days'")

            rows = cur.fetchall()

            for row in rows:

                if row != None:

                    created_at_lst.append(row[0])

                    id_lst.append(row[1])

                    email_lst.append(row[2])

                    ip_lst.append(row[3])

            cur.close()

        except (Exception, psycopg2.DatabaseError) as error:

            print (error)

        finally:

            if conn is not None:

                conn.close()

        return (created_at_lst, id_lst, email_lst, ip_lst)

    def save_registers(self, created_at_lst, id_lst, email_lst, ip_lst):

        insert_sql = 'INSERT INTO spamcheck(created_at, id, email, ip, tor_exit_node) VALUES(%s,%s,%s,%s,%s) ON CONFLICT DO NOTHING'
        
        i = 0

        while i < len(id_lst):

            is_tor_exit_node = self.__check_ip(self, ip_lst[i])

            tor_exit_node = 't' if is_tor_exit_node == 't' else 'f'

            conn = None

            try:

                conn = psycopg2.connect(database = self.spamcheck_db, user = self.spamcheck_db_user, password = "", host = "/var/run/postgresql", port = "5432")

                cur = conn.cursor()

                cur.execute(insert_sql, (created_at_lst[i], id_lst[i], email_lst[i], ip_lst[i], tor_exit_node))

                conn.commit()

                cur.close()

            except (Exception, psycopg2.DatabaseError) as error:

                print(error)

            finally:

                if conn is not None:

                    conn.close()
            
            print(created_at_lst[i], id_lst[i], email_lst[i], ip_lst[i], tor_exit_node)

            i += 1

    def get_totals(self):

        spamcheck_datetime_lst = []

        spamcheck_registers_lst = []

        select_sql = 'select date(created_at), count(ip) as registers from spamcheck group by date(created_at) order by date(created_at)'

        conn = None

        try:

            conn = psycopg2.connect(database = self.spamcheck_db, user = self.spamcheck_db_user, password = "", host = "/var/run/postgresql", port = "5432")

            cur = conn.cursor()

            cur.execute(select_sql)

            rows = cur.fetchall()

            for row in rows:

                spamcheck_datetime_lst.append(row[0])

                spamcheck_registers_lst.append(row[1])

            cur.close()

        except (Exception, psycopg2.DatabaseError) as error:

            print(error)

        finally:

            if conn is not None:

                conn.close()

        return (spamcheck_datetime_lst, spamcheck_registers_lst)

    def write_totals(self, spamcheck_datetime_lst, spamcheck_registers_lst):

        insert_sql = 'INSERT INTO totals(datetime, registers) VALUES(%s,%s) ON CONFLICT (datetime) DO UPDATE SET (datetime, registers) = (EXCLUDED.datetime, EXCLUDED.registers)'

        first_date = spamcheck_datetime_lst[0]

        last_date = spamcheck_datetime_lst[len(spamcheck_datetime_lst)-1]

        i = 0

        while i < len(spamcheck_datetime_lst):

            conn = None

            try:

                conn = psycopg2.connect(database = self.spamcheck_db, user = self.spamcheck_db_user, password = "", host = "/var/run/postgresql", port = "5432")

                cur = conn.cursor()

                if first_date == spamcheck_datetime_lst[i]:

                    cur.execute(insert_sql, (spamcheck_datetime_lst[i], spamcheck_registers_lst[i]))

                    i += 1

                else:

                    cur.execute(insert_sql, (first_date, '0'))

                conn.commit()

                cur.close()

            except (Exception, psycopg2.DatabaseError) as error:

                print(error)

            finally:

                if conn is not None:

                    conn.close()

            first_date = first_date + timedelta(days=1)

        if date.today() == last_date + timedelta(days=1):

            insert_sql = 'INSERT INTO totals(datetime, registers) VALUES(%s,%s) ON CONFLICT (datetime) DO UPDATE SET (datetime, registers) = (EXCLUDED.datetime, EXCLUDED.registers)'

            conn = None

            try:

                conn = psycopg2.connect(database = self.spamcheck_db, user = self.spamcheck_db_user, password = "", host = "/var/run/postgresql", port = "5432")

                cur = conn.cursor()

                cur.execute(insert_sql, (date.today(), '0'))

                conn.commit()

                cur.close()

            except (Exception, psycopg2.DatabaseError) as error:

                print(error)

            finally:

                if conn is not None:

                    conn.close()

    def check_approval(self, user_id):

        approved = False

        try:

            conn = None

            conn = psycopg2.connect(database = self.mastodon_db, user = self.mastodon_db_user, password = "", host = "/var/run/postgresql", port = "5432")

            cur = conn.cursor()

            cur.execute("select approved from users where id = (%s)", (user_id,))

            row = cur.fetchone()

            if row != None:

                approved = row[0]

            cur.close()

            return approved

        except (Exception, psycopg2.DatabaseError) as error:

            print (error)

        finally:

            if conn is not None:

                conn.close()

    @staticmethod
    def __check_ip(self, ip):

        is_tor_exit_node = 'f'

        if ip == None:

            return

        conn = None

        try:

            conn = psycopg2.connect(database = self.spamcheck_db, user = self.spamcheck_db_user, password = "", host = "/var/run/postgresql", port = "5432")

            cur = conn.cursor()

            cur.execute('select ip from torexit_ips where ip=(%s)', (ip,))

            row = cur.fetchone()

            if row != None:

                is_tor_exit_node = 't'

            cur.close()

        except (Exception, psycopg2.DatabaseError) as error:

            print(error)

        finally:

            if conn is not None:

                conn.close()

        return is_tor_exit_node

    @staticmethod
    def __check_setup(self):

        is_setup = False

        if not os.path.isfile(self.config_file):
            print(f"File {self.config_file} not found, running setup.\n")
        else:
            is_setup = True

        return is_setup
    
    @staticmethod
    def __setup(self):

        if not os.path.exists('config'):
            os.makedirs('config')

        self.mastodon_hostname = input("Mastodon hostname, in ex. 'mastodon.social': ")
        self.mastodon_db = input("Mastodon's database name: ")
        self.mastodon_db_user = input("Mastodon's database user: ")
        self.spamcheck_db = input("Spamcheck's database name: ")
        self.spamcheck_db_user = input("Spamcheck's database user: ")

        if not os.path.exists(self.config_file):
            with open(self.config_file, 'w'): pass
            print(f"\n{self.config_file} created!\n")

        with open(self.config_file, 'a') as the_file:
            print(f"Writing Mastodon hostname parameter to {self.config_file}")
            the_file.write(f'mastodon_hostname: {self.mastodon_hostname}\n')
            the_file.write(f'mastodon_db: {self.mastodon_db}\n')
            the_file.write(f'mastodon_db_user: {self.mastodon_db_user}\n')
            the_file.write(f'spamcheck_db: {self.spamcheck_db}\n')
            the_file.write(f'spamcheck_db_user: {self.spamcheck_db_user}\n')

        return (self.mastodon_hostname, self.mastodon_db, self.mastodon_db_user, self.spamcheck_db, self.spamcheck_db_user)

    @staticmethod
    def __check_dbsetup(self):

        dbsetup = False

        try:

            conn = None

            conn = psycopg2.connect(database = self.spamcheck_db, user = self.spamcheck_db_user, password = "", host = "/var/run/postgresql", port = "5432")

            dbsetup = True

        except (Exception, psycopg2.DatabaseError) as error:

            print(error)

        return dbsetup

    @staticmethod
    def __createdb(self):

        conn = None

        try:

            conn = psycopg2.connect(dbname='postgres',
                user=self.spamcheck_db_user, host='',
                password='')

            conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)

            cur = conn.cursor()

            print(f"Creating database {self.spamcheck_db}. Please wait...")

            cur.execute(sql.SQL("CREATE DATABASE {}").format(
                sql.Identifier(self.spamcheck_db))
            )
            print(f"Database {self.spamcheck_db} created!\n")

            self.__dbtables_schemes(self)

        except (Exception, psycopg2.DatabaseError) as error:

            print(error)

        finally:

            if conn is not None:

                conn.close()

    @staticmethod
    def __dbtables_schemes(self):

        table = "spamcheck"
        sql = "create table "+table+" (created_at timestamptz, id bigint PRIMARY KEY, email varchar(200), ip inet, tor_exit_node boolean)"
        self.__create_table(self, table, sql)

        table = "torexit_ips"
        sql = "create table "+table+" (created_at timestamptz, ip inet PRIMARY KEY)"
        self.__create_table(self, table, sql)

        table = "totals"
        sql = "create table "+table+" (datetime timestamptz PRIMARY KEY, registers int)"
        self.__create_table(self, table, sql)

    @staticmethod
    def __create_table(self, table, sql):

        conn = None

        try:

            conn = psycopg2.connect(database = self.spamcheck_db, user = self.spamcheck_db_user, password = "", host = "/var/run/postgresql", port = "5432")
            cur = conn.cursor()

            print(f"Creating table {table}")
            cur.execute(sql)

            conn.commit()
            print(f"Table {table} created!\n")

        except (Exception, psycopg2.DatabaseError) as error:

            print(error)

        finally:

            if conn is not None:

                conn.close()

    @staticmethod
    def __get_parameter(parameter, file_path ):

        with open( file_path ) as f:
            for line in f:
                if line.startswith( parameter ):
                    return line.replace(parameter + ":", "").strip()

        print(f'{file_path} Missing parameter {parameter}')
        sys.exit(0)

###############################################################################
# main

if __name__ == '__main__':

    spamcheck = Spamcheck()
    
    created_at_lst, id_lst, email_lst, ip_lst = spamcheck.new_registers()

    spamcheck.save_registers(created_at_lst, id_lst, email_lst, ip_lst)

    spamcheck_datetime_lst, spamcheck_registers_lst = spamcheck.get_totals()

    spamcheck.write_totals(spamcheck_datetime_lst, spamcheck_registers_lst)