import os
import sys
import psycopg2
import pdb

def check_ip(ip):

    is_tor_exit_node = 'f'

    conn = None

    try:

        conn = psycopg2.connect(database = spamcheck_db, user = 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:

            print(f'{ip} is a Tor exit node')

            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

def get_spam_ips():

    spam_ip_lst = []

    conn = None

    try:

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

        cur = conn.cursor()

        cur.execute('select ip from spamcheck')

        rows = cur.fetchall()

        for row in rows:

            spam_ip_lst.append(row[0])

        cur.close()

    except (Exception, psycopg2.DatabaseError) as error:

        print(error)

    finally:

        if conn is not None:

            conn.close()

    return spam_ip_lst

def db_config():

    # Load db configuration from config file
    config_filepath = "config/db_config.txt"
    spamcheck_db = get_parameter("spamcheck_db", config_filepath)
    spamcheck_db_user = get_parameter("spamcheck_db_user", config_filepath)

    return (spamcheck_db, spamcheck_db_user)

# Returns the parameter from the specified file
def get_parameter( parameter, file_path ):
    # Check if secrets file exists
    if not os.path.isfile(file_path):
        print("File %s not found, exiting."%file_path)
        sys.exit(0)

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

    # Cannot find parameter, exit
    print(file_path + "  Missing parameter %s "%parameter)
    sys.exit(0)

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

if __name__ == '__main__':

    spamcheck_db, spamcheck_db_user = db_config()

    spam_ip_lst = get_spam_ips()

    print(f'{len(spam_ip_lst)} IPs found.')

    i = 0

    while i < len(spam_ip_lst):

        is_tor_exit_node = check_ip(spam_ip_lst[i])

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

        update_sql = 'UPDATE spamcheck set tor_exit_node=(%s) where ip=(%s)'

        conn = None

        try:

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

            cur = conn.cursor()

            cur.execute(update_sql, (tor_exit_node, spam_ip_lst[i]))

            conn.commit()

            cur.close()

        except (Exception, psycopg2.DatabaseError) as error:

            print(error)

        finally:

            if conn is not None:

                conn.close()

        i += 1