budget/budget.py

import os
from datetime import datetime, timezone, timedelta, date
import time
import pytz
import sys
import psycopg2
from mastodon import Mastodon
from mastodon.Mastodon import MastodonMalformedEventError, MastodonNetworkError, MastodonReadTimeout, MastodonAPIError, MastodonIllegalArgumentError
from prettytable import PrettyTable
import pdb

tz = pytz.timezone('Europe/Madrid')

menu_options = {
    1: 'Add Bill',
    2: 'List Bills',
    3: 'Add Donation',
    4: 'List Donations',
    5: 'Balance',
    6: 'Exit',
}
list_show_options = {
    1: 'Month',
    2: 'Year',
    3: 'Total',
    4: 'Back',
}

def print_menu():

    for key in menu_options.keys():
        print (key, '-', menu_options[key] )

def list_show_menu():

    for key in list_show_options.keys():
        print (key, '-', list_show_options[key] )

def check_db_conn():

    try:

        conn = None

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

    except (Exception, psycopg2.DatabaseError) as error:

        sys.stdout.write(f'\n{str(error)}\n')

        sys.exit("Exiting. Run 'python db-setup'")

class Bill:

    name = "Bill"

    def __init__(self):

        self.domain = float(input('Enter Domain Bill: ') or '0.00')
        self.server = float(input('Server Bill: ') or '0.00')
        self.backup = float(input('Backup Bill: ') or '0.00')
        self.fileserver = float(input("File server Bill: ") or '0.00')
        self.setup = float(input(f'Setup Bill: ') or '0.00')
        
        while True:

            try:

                self.date = datetime.strptime(input('Enter Bill date in the format d.m.yyyy: '), '%d.%m.%Y')

                now = datetime.now(tz)

                self.date = self.date.replace(hour=now.hour, minute=now.minute, second=now.second)

                break

            except:

                self.date = format(datetime.now(tz).strftime('%d.%m.%Y %H:%M:%S'))

                break

    def save(self):

        sql  = "INSERT INTO bills(datetime, domain, server, backup, fileserver, setup) VALUES(%s,%s,%s,%s,%s,%s)"

        try:

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

            cur = conn.cursor()

            cur.execute(sql, (self.date, self.domain, self.server, self.backup, self.fileserver, self.setup))
            
            print("\nUpdating bills...")

            conn.commit()

            cur.close()

        except (Exception, psycopg2.DatabaseError) as error:

            print(error)

        finally:

            if conn is not None:

                conn.close()

class Donation:

    name = "Donation"

    def __init__(self):

        self.amount = float(input('Donation: ') or '0.00')

        while True:

            try:

                self.date = datetime.strptime(input('Enter Donation date in the format d.m.yyyy: '), '%d.%m.%Y')

                now = datetime.now(tz)

                self.date = self.date.replace(hour=now.hour, minute=now.minute, second=now.second)

                break

            except:

                self.date = format(datetime.now(tz).strftime('%d.%m.%Y %H:%M:%S'))

                break

    def save(self):

        sql  = "INSERT INTO incomes(datetime, donations) VALUES(%s,%s)"

        try:

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

            cur = conn.cursor()

            cur.execute(sql, (self.date, self.amount))

            print("\nUpdating donations...")

            conn.commit()

            cur.close()

        except (Exception, psycopg2.DatabaseError) as error:

            print(error)

        finally:

            if conn is not None:

                conn.close()

class BillShow:

    name = "BillShow"

    def __init__(self):

        while(True):

            print('\n')

            list_show_menu()

            show_option = ''

            now = datetime.now(tz)

            try:

                show_option = int(input('\nEnter your choice: '))

            except:

                print('\nWrong input. Please enter a number between 1 and 4.\n')

            if show_option == 1:

                self.option = 1
                self.criteria = input('Month (enter = current month)? ') or str(now.month)
                break

            elif show_option == 2:

                self.option = 2
                self.criteria = input('Year (enter = current year)? ') or str(now.year)
                break

            elif show_option == 3:

                self.option = 3
                self.criteria = None
                break

            elif show_option == 4:

                self.option = 4
                self.criteria = None
                break

    def show(self, criteria):

        month_date = []

        month_domain = []

        month_server = []

        month_backup = []

        month_fileserver = []

        month_setup = []

        if self.option == 1:

            sql  = "select datetime, domain, server, backup, fileserver, setup from bills where date_part('year', datetime) = date_part('year', CURRENT_DATE) and date_part('month', datetime) = (%s) order by 1 asc"

        elif self.option == 2:

            sql  = "select datetime, domain, server, backup, fileserver, setup from bills where date_part('year', datetime) = (%s) order by 1 asc"

        elif self.option == 3:

            sql  = "select datetime, domain, server, backup, fileserver, setup from bills order by 1 asc"

        try:

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

            cur = conn.cursor()

            if self.option == 1 or self.option == 2:

                cur.execute(sql, (self.criteria,))

            else:
                
                cur.execute(sql)

            rows = cur.fetchall()

            for row in rows:

                month_date.append(row[0])

                month_domain.append(row[1])

                month_server.append(row[2])

                month_backup.append(row[3])

                month_fileserver.append(row[4])

                month_setup.append(row[5])
            
            cur.close()

        except (Exception, psycopg2.DatabaseError) as error:

            print(error)

        finally:

            if conn is not None:

                conn.close()

        return (month_date, month_domain, month_server, month_backup, month_fileserver, month_setup)
    
class DonationShow:

    name = "DonationShow"

    def __init__(self):

        while(True):

            print('\n')

            list_show_menu()

            show_option = ''

            now = datetime.now(tz)

            try:

                show_option = int(input('\nEnter your choice: '))

            except:

                print('\nWrong input. Please enter a number between 1 and 4.\n')

            if show_option == 1:

                self.option = 1
                self.criteria = input('Month (enter = current month)? ') or str(now.month)
                break

            elif show_option == 2:

                self.option = 2
                self.criteria = input('Year (enter = current year)? ') or str(now.year)
                break

            elif show_option == 3:

                self.option = 3
                self.criteria = None
                break

            elif show_option == 4:

                self.option = 4
                self.criteria = None
                break

    def show(self, criteria):

        month_date = []

        month_donation = []

        if self.option == 1:

            sql  = "select datetime, donations from incomes where date_part('year', datetime) = date_part('year', CURRENT_DATE) and date_part('month', datetime) = (%s) order by 1 asc"

        elif self.option == 2:

            sql  = "select datetime, donations from incomes where date_part('year', datetime) = (%s) order by 1 asc"

        elif self.option == 3:

            sql  = "select datetime, donations from incomes order by 1 asc"

        try:

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

            cur = conn.cursor()

            if self.option == 1 or self.option == 2:

                cur.execute(sql, (self.criteria,))

            else:

                cur.execute(sql)

            rows = cur.fetchall()

            for row in rows:

                month_date.append(row[0])

                month_donation.append(row[1])

            cur.close()

        except (Exception, psycopg2.DatabaseError) as error:

            print(error)

        finally:

            if conn is not None:

                conn.close()

        return (month_date, month_donation)

class BalanceShow:

    name = "BalanceShow"

    def __init__(self):

        while(True):

            print('\n')

            list_show_menu()

            show_option = ''

            now = datetime.now(tz)

            try:

                show_option = int(input('\nEnter your choice: '))

            except:

                print('\nWrong input. Please enter a number between 1 and 4.\n')

            if show_option == 1:

                self.option = 1
                self.criteria = input('Month (enter = current month)? ') or str(now.month)
                break

            elif show_option == 2:

                self.option = 2
                self.criteria = input('Year (enter = current year)? ') or str(now.year)
                break

            elif show_option == 3:

                self.option = 3
                self.criteria = None
                break

            elif show_option == 4:

                self.option = 4
                self.criteria = None
                break

    def show(self, criteria):

        balance = 0

        if self.option == 1:

            incomes_sql = "select sum(donations) from incomes where date_part('year', datetime) = date_part('year', CURRENT_DATE) and date_part('month', datetime) = (%s)"

            bills_sql  = "select sum(coalesce(domain) + coalesce(server) + coalesce(backup) + coalesce(fileserver) + coalesce(setup)) from bills where date_part('year', datetime) = date_part('year', CURRENT_DATE) and date_part('month', datetime) = (%s)"

        elif self.option == 2:

            incomes_sql  = "select sum(donations) from incomes where date_part('year', datetime) = (%s)"

            bills_sql  = "select sum(coalesce(domain) + coalesce(server) + coalesce(backup) + coalesce(fileserver) + coalesce(setup)) from bills where date_part('year', datetime) = (%s)"

        elif self.option == 3:

            incomes_sql  = "select sum(donations) from incomes order by 1 asc"

            bills_sql = "select sum(coalesce(domain) + coalesce(server) + coalesce(backup) + coalesce(fileserver) + coalesce(setup)) from bills"

        try:

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

            cur = conn.cursor()

            if self.option == 1 or self.option == 2:

                cur.execute(incomes_sql, (self.criteria,))

                row = cur.fetchone()

                incomes = row[0] if row[0] != None else 0

                cur.execute(bills_sql, (self.criteria,))

                row = cur.fetchone()

                bills = row[0]

                bills = row[0] if row[0] != None else 0

            else:

                cur.execute(incomes_sql)

                row = cur.fetchone()

                incomes = row[0]

                cur.execute(bills_sql)

                row = cur.fetchone()

                bills = row[0]

            balance = incomes - bills

            cur.close()

        except (Exception, psycopg2.DatabaseError) as error:

            print(error)

        finally:

            if conn is not None:

                conn.close()

        return (balance)

def get_config():

    global budget_db
    global budget_db_user

    # Load configuration from config file
    config_filepath = "config/db_config.txt"
    budget_db = get_db_params("budget_db", config_filepath)
    budget_db_user = get_db_params("budget_db_user", config_filepath)

    return (config_filepath, budget_db, budget_db_user)

def get_db_params( parameter, file_path ):

    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__':

    get_config()

    check_db_conn()

    while(True):

        print('\n')

        print_menu()

        option = ''

        try:

            option = int(input('\nEnter your choice: '))

        except:

            print('\nWrong input. Please enter a number between 1 and 5.\n')

        if option == 1:

            bill = Bill()

            if bill.domain == bill.server == bill.backup == bill.fileserver == bill.setup == 0.0:

                print('\nAny bill added')

            else:

                bill.save()

        elif option == 2:

            bill = BillShow()

            if bill.option == 1 or bill.option == 2 or bill.option == 3:

                month_date, month_domain, month_server, month_backup, month_fileserver, month_setup = bill.show(bill.criteria)
            
            if bill.option != 4:

                print_table = PrettyTable()

                print_table.field_names = ['Date', 'Domain', 'Server', 'Backup', 'File server', 'Setup']

                i = 0

                total_amount = 0

                while i < len(month_date):

                    b_date = format(month_date[i].strftime('%d.%m.%Y %H:%M:%S'))

                    b_domain = float(month_domain[i])

                    b_server = float(month_server[i])

                    b_backup = float(month_backup[i])

                    b_fileserver = float(month_fileserver[i])

                    b_setup = float(month_setup[i])

                    total_amount = total_amount + b_domain + b_server + b_backup + b_fileserver + b_setup

                    print_table.add_row([b_date, b_domain, b_server, b_backup, b_fileserver, b_setup])

                    i += 1

                print_table.add_row(['', '', '', '', 'Total', round(total_amount, 2)])

                print(print_table)

        elif option == 3:

            donation = Donation()

            if donation.amount == 0.0:

                print('\nAny donation added')

            else:

                donation.save()

        elif option == 4:

            donation = DonationShow()

            if donation.option == 1 or donation.option == 2 or donation.option == 3: 
            
                date_lst, donation_lst = donation.show(donation.criteria)

            if donation.option != 4:

                print_table = PrettyTable()

                print_table.field_names = ['Date','Donation']

                print_table.align['Donation'] = "r"

                i = 0

                total_amount = 0

                while i < len(donation_lst):

                    d_date = format(date_lst[i].strftime('%d.%m.%Y %H:%M:%S'))

                    d_amount = float(donation_lst[i])

                    total_amount = total_amount + d_amount

                    print_table.add_row([d_date, d_amount])

                    i += 1

                print_table.add_row(['Total', round(total_amount, 2)])

                print(print_table)

        elif option == 5:

            balance = BalanceShow()

            if balance.option == 1 or balance.option == 2 or balance.option == 3: 

                balance_result  = balance.show(balance.criteria)

            if balance.option != 4:

                print_table = PrettyTable()

                print_table.field_names = ['Balance']

                print_table.align['Balance'] = "r"

                b_amount = float(balance_result)

                print_table.add_row([b_amount])

                #print_table.add_row(['Total', round(total_amount, 2)])

                print(print_table)

        elif option == 6:

            print('Bye!')

            exit()