diff --git a/commands.py b/commands.py index 2adad376bb725512e69fece4ac3d9f8ea4d314e6..bb95c05f8b63a85f86d53996251c7c830bf0fab2 100644 --- a/commands.py +++ b/commands.py @@ -1,88 +1,139 @@ +import ast import os -from bitcoinlib.keys import HDKey from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup from telegram.ext import CallbackContext -from bitcoinlib.wallets import Wallet -from models.AppWallet import AppWallet +from database import AppUser, AppState, AppWallet, AppWalletRequest from util import DBManager -from models.AppUser import AppUser -dbmanager = DBManager() +# dbmanager = DBManager() + +def generate_message(state_id): + state = AppState.get(AppState.id == state_id) + reply_msg = state.message + menu_string = ast.literal_eval(state.menu) + menu = [[InlineKeyboardButton(y[0], callback_data=y[1]) for y in x] for x in menu_string] + reply_markup = InlineKeyboardMarkup(menu) + return reply_msg, reply_markup def start(update: Update, context: CallbackContext): - user = dbmanager.create_or_get_user(update.message.from_user) - menu = [[InlineKeyboardButton('Create a wallet', callback_data='create_a_wallet'), InlineKeyboardButton('Send a transaction', callback_data='send_a_transaction')]] - reply_markup = InlineKeyboardMarkup(menu) + telegram_user = update.effective_user + user, created = AppUser.get_or_create(telegram_id=telegram_user.id, nickname=telegram_user.full_name) + user.set_state(1) + # TODO: clear variables + + reply_msg, reply_markup = generate_message(1) + update.effective_message.reply_text(reply_msg, reply_markup=reply_markup) + - update.message.reply_text("Hi! Welcome! What do you want to do?", reply_markup=reply_markup) +def back(update: Update, context: CallbackContext): + telegram_user = update.effective_user + user, created = AppUser.get_or_create(telegram_id=telegram_user.id, nickname=telegram_user.full_name) + state = AppState.get(AppState.id == user.state_id) + user.set_state(state.prev_state) + + reply_msg, reply_markup = generate_message(state.prev_state) + update.effective_message.reply_text(reply_msg, reply_markup=reply_markup) def send_a_transaction(update: Update, context: CallbackContext): - user = dbmanager.create_or_get_user(update.message.from_user) + telegram_user = update.effective_user + user, created = AppUser.get_or_create(telegram_id=telegram_user.id, nickname=telegram_user.full_name) pass def create_wallet(update: Update, context: CallbackContext): - user = dbmanager.create_or_get_user(update.message.from_user) - user.set_state(1) - update.message.reply_text("OK, enter co-signers number") - pass + telegram_user = update.effective_user + user, created = AppUser.get_or_create(telegram_id=telegram_user.id, nickname=telegram_user.full_name) + user.set_state(2) + reply_msg, reply_markup = generate_message(2) + update.effective_message.reply_text(reply_msg, reply_markup=reply_markup) -def enter_co_signers(update: Update, context: CallbackContext): - user = dbmanager.create_or_get_user(update.message.from_user) - msg = update.message.text +def enter_co_signers(update: Update, context: CallbackContext, user: AppUser): + msg = update.effective_message.text try: co_signers = int(msg) - user.variable['co_signers'] = co_signers - user.set_state(2) - update.message.reply_text("OK, enter the minimum signature number") + user.set_variable('max_co_signers', co_signers) + next_state = user.next_state() + reply_msg, reply_markup = generate_message(next_state) + update.effective_message.reply_text(reply_msg, reply_markup=reply_markup) except ValueError: - update.message.reply_text("This is not a number, try again") + update.effective_message.reply_text("This is not a number, try again") -def minimum_co_signers(update: Update, context: CallbackContext): - user = dbmanager.create_or_get_user(update.message.from_user) - msg = update.message.text +def minimum_co_signers(update: Update, context: CallbackContext, user: AppUser): + msg = update.effective_message.text try: min_signers = int(msg) - user.variable['min_signers'] = min_signers - user.set_state(3) - update.message.reply_text("Great! Now choose a name for this wallet") + user.set_variable('min_co_signers', min_signers) + next_state = user.next_state() + reply_msg, reply_markup = generate_message(next_state) + update.effective_message.reply_text(reply_msg, reply_markup=reply_markup) except ValueError: - update.message.reply_text("This is not a number, try again") + update.effective_message.reply_text("This is not a number, try again") -def name_wallet(update: Update, context: CallbackContext): - user = dbmanager.create_or_get_user(update.message.from_user) - msg = update.message.text +def name_wallet(update: Update, context: CallbackContext, user: AppUser): + msg = update.effective_message.text if len(msg) > 100: - update.message.reply_text("The chosen name is too long, please try again") + update.effective_message.reply_text("The chosen name is too long, please try again") else: - user.variable['wallet_name'] = msg - AppWallet.create_wallet(user.variable['co_signers'], - user.variable['min_signers'], - user.id, - user.variable['wallet_name']) + user.set_variable('wallet_name', msg) + variables = user.all_variables + app_wallet = AppWallet.create(max_co_signers=variables['max_co_signers'], + min_co_signers=variables['min_co_signers'], + initiator_user_id=user.id, + name=variables['wallet_name']) + + wallet_requests = AppWalletRequest.generate_requests(variables['max_co_signers'] - 1, app_wallet.id) + update.effective_message.reply_text("Wallet is created, forward each of these messages to one co-signer") + + for request in wallet_requests: + update.effective_message.reply_text("Hi, please join this wallet as co-signer.\n" + "Wallet name= {}\n" + "https://t.me/BitcoinLibTestbot$start={}".format(variables['wallet_name'], request["token"])) + + next_state = user.next_state() + reply_msg, reply_markup = generate_message(next_state) + update.effective_message.reply_text(reply_msg, reply_markup=reply_markup) + + +def callback_handler(update: Update, context: CallbackContext): + query = update.callback_query.data + try: + f = COMMANDS[query] + f(update, context) + except KeyError: + # Current state does not have a message handler, so do nothing + pass def message_handler(update: Update, context: CallbackContext): - user = dbmanager.create_or_get_user(update.message.from_user) - state_user = user.state + telegram_user = update.effective_user + user, created = AppUser.get_or_create(telegram_id=telegram_user.id, nickname=telegram_user.full_name) + state_user = user.state_id.id HANDLERS = { - 1: enter_co_signers, - 2: minimum_co_signers, - 3: name_wallet + 2: enter_co_signers, + 3: minimum_co_signers, + 4: name_wallet } try: f = HANDLERS[state_user] - f(update, context) + f(update, context, user) except KeyError: # Current state does not have a message handler, so do nothing pass + + +COMMANDS = { + "start": start, + "create_a_wallet": create_wallet, + "send_a_transaction": send_a_transaction, + "back": back, +} diff --git a/database.py b/database.py index 9f86a6d3773629e2ca5efb797998e00463969465..d217041e3b9678d73deff7a9fc8985bb7875f90c 100644 --- a/database.py +++ b/database.py @@ -1,9 +1,14 @@ -from peewee import MySQLDatabase, Model, CharField, BlobField, AutoField +import json +import random +import string + +from peewee import MySQLDatabase, Model, CharField, BlobField, AutoField, IntegrityError, DoesNotExist, BigIntegerField from peewee import IntegerField, ForeignKeyField, DoubleField import os db = MySQLDatabase(None) + def connect(create_tables=False) -> None: """Initializes the database session and connects to the database :param create_tables: Creates database tables [default: False] @@ -17,47 +22,77 @@ def connect(create_tables=False) -> None: db.connect() if create_tables: - db.create_tables([State, User, Wallet, - WalletRequest, UserWallet, PendingTX]) + db.create_tables([AppState, AppUser, AppWallet, + AppWalletRequest, AppUserWallet, AppPendingTX]) + class BaseModel(Model): id = AutoField() + class Meta: database = db -class State(BaseModel): +class AppState(BaseModel): message = CharField() - menu = CharField() + menu = CharField(null=True) prev_state = IntegerField() + next_state = IntegerField() -class User(BaseModel): - telegram_id = IntegerField() - state_id = ForeignKeyField(State) +class AppUser(BaseModel): + telegram_id = BigIntegerField(unique=True) + state_id = ForeignKeyField(AppState, default=1) nickname = CharField() - variables = CharField() + variables = CharField(default="{}") + + def next_state(self): + AppUser.update({AppUser.state_id: self.state_id.next_state}).where(AppUser.id == self.id).execute() + return self.state_id.next_state + + def set_state(self, new_state): + AppUser.update({AppUser.state_id: new_state}).where(AppUser.id == self.id).execute() + @property + def all_variables(self): + newuser = AppUser.get(AppUser.id == self.id) + return json.loads(str(newuser.variables)) -class Wallet(BaseModel): + def set_variables(self, new_variables: dict): + AppUser.update({AppUser.variables: json.dumps(new_variables)}).where(AppUser.id == self.id).execute() + + def set_variable(self,key , value): + new_variables = self.all_variables + new_variables[key] = value + self.set_variables(new_variables) + + +class AppWallet(BaseModel): max_co_signers = IntegerField() min_co_signers = IntegerField() - initiator_user_id = ForeignKeyField(User, backref='users') + initiator_user_id = ForeignKeyField(AppUser, backref='users') name = CharField() -class WalletRequest(BaseModel): +class AppWalletRequest(BaseModel): token = CharField() - wallet_id = ForeignKeyField(User, backref='walletrequests') + wallet_id = ForeignKeyField(AppWallet, backref='walletrequests') + + @classmethod + def generate_requests(cls, amount, wallet_id): + wallet_requests = [{'token': ''.join(random.choice(string.ascii_letters) for i in range(64)), 'wallet_id': wallet_id} for x in range(amount)] + AppWalletRequest.insert(wallet_requests).execute() + return wallet_requests + -class UserWallet(BaseModel): - user_id = ForeignKeyField(User, backref='wallets') - wallet_id = ForeignKeyField(Wallet, backref='wallets') +class AppUserWallet(BaseModel): + user_id = ForeignKeyField(AppUser, backref='wallets') + wallet_id = ForeignKeyField(AppWallet, backref='wallets') -class PendingTX(BaseModel): - from_wallet_id = ForeignKeyField(Wallet, backref='pendingtx') +class AppPendingTX(BaseModel): + from_wallet_id = ForeignKeyField(AppWallet, backref='pendingtx') to_address = CharField() amount = DoubleField() fee_sat_per_byte = IntegerField() diff --git a/main.py b/main.py index 0d4ad500b11af451f85db9adc34458f821d3e08f..3fa963ec7eec6832bbd61b3a4138f907f06f33dc 100644 --- a/main.py +++ b/main.py @@ -1,7 +1,8 @@ import os from dotenv import load_dotenv -from telegram.ext import Updater, CommandHandler, MessageHandler, Filters +from telegram.ext import Updater, CommandHandler, MessageHandler, Filters, CallbackQueryHandler import commands +from database import connect load_dotenv() @@ -13,16 +14,20 @@ class TelegramWrapper: self._updater = Updater(token=token) self._dispatcher = self._updater.dispatcher self.registerCommands() + self.connect_database() def registerCommands(self): - self._dispatcher.add_handler(CommandHandler('start', commands.start)) - self._dispatcher.add_handler(CommandHandler('create_a_wallet', commands.create_wallet)) - self._dispatcher.add_handler(CommandHandler('send_a_transaction', commands.send_a_transaction)) + for key in commands.COMMANDS: + self._dispatcher.add_handler(CommandHandler(key, commands.COMMANDS[key])) + self._dispatcher.add_handler(CallbackQueryHandler(commands.callback_handler)) self._dispatcher.add_handler(MessageHandler(Filters.all, commands.message_handler)) def start_polling(self): self._updater.start_polling() + def connect_database(self): + connect(True) + if __name__ == "__main__": telegramWrapper = TelegramWrapper(os.getenv("TOKEN")) diff --git a/models/AppUser.py b/models/AppUser.py deleted file mode 100644 index df86eaaa0cb2209763246fc85c294d9ed794b39f..0000000000000000000000000000000000000000 --- a/models/AppUser.py +++ /dev/null @@ -1,9 +0,0 @@ -class AppUser: - - def __init__(self, state): - self.state = state - pass - - def set_state(self, new_state): - self.state = new_state - # TODO: change it in the database diff --git a/models/AppUserWallet.py b/models/AppUserWallet.py deleted file mode 100644 index 16753a65c6415da4178437f50a5b4b15094ff58c..0000000000000000000000000000000000000000 --- a/models/AppUserWallet.py +++ /dev/null @@ -1,10 +0,0 @@ -class AppUserWallet: - - def __init__(self, user_id, wallet_id): - self.user_id = user_id - self.wallet_id = wallet_id - - @classmethod - def create(cls, user_id, wallet_id): - user_wallet = UserWallet(user_id, wallet_id) - # TODO: PUT THIS OBJECT INTO THE DATABASE diff --git a/models/AppWallet.py b/models/AppWallet.py deleted file mode 100644 index e2ec89c18a37b047e2c77eb683bb0a025a0de64e..0000000000000000000000000000000000000000 --- a/models/AppWallet.py +++ /dev/null @@ -1,9 +0,0 @@ - -class AppWallet: - def __init__(self): - pass - - @classmethod - def create_wallet(cls,max_co_signers, min_co_signers, initiator_id, wallet_name): - # TODO create new wallet in the database - pass