Added testing and premium accounts and commands

wallamanta/helpers.py

@@ -1,3 +1,5 @@
+from datetime import datetime, timedelta
+
 #TELEGRAM_ESCAPE_CHARACTERS = ['_', '*', '[', ']', '(', ')', '~', '>', '#', '+', '-', '=', '|', '{', '}', '.', '!']
 TELEGRAM_ESCAPE_CHARACTERS = ['_', '*', '[', ']', '(', ')', '~', '>', '+', '-', '=', '|', '{', '}', '.', '!']
 ADMIN_IDS = [10101691]
@@ -15,3 +17,7 @@ def is_user_admin(telegram_user_id):
 
 def get_telegram_user_id(update):
     return update.message.chat_id
+
+def get_date_ahead(add_days):
+    date_ahead = datetime.today() + timedelta(days=add_days)
+    return date_ahead.strftime("%d/%m/%Y")

wallamanta/walladb.py

@@ -23,7 +23,7 @@ def dict_factory(cursor, row):
 def setup_db():
     con = sqlite3.connect(DB)
     cur = con.cursor()
-    cur.execute("CREATE TABLE IF NOT EXISTS users(telegram_user_id, active)")
+    cur.execute("CREATE TABLE IF NOT EXISTS users(telegram_user_id, active, type, until)")
     cur.execute("CREATE TABLE IF NOT EXISTS products(product_name, distance, \
                 latitude, longitude, condition, min_price, max_price, \
                 title_exclude, title_description_exclude, telegram_user_id)")
@@ -37,21 +37,50 @@ def is_user_valid(telegram_user_id):
     con.close()
     return ret
 
-def add_valid_user(telegram_user_id):
+def is_user_premium(telegram_user_id):
+    con = sqlite3.connect(DB)
+    cur = con.cursor()
+    res = cur.execute(f"SELECT * FROM users WHERE telegram_user_id={telegram_user_id} AND active=True AND type='premium'")
+    ret = res.fetchone() != None
+    con.close()
+    return ret
+
+def is_user_testing(telegram_user_id):
+    con = sqlite3.connect(DB)
+    cur = con.cursor()
+    res = cur.execute(f"SELECT * FROM users WHERE telegram_user_id={telegram_user_id} AND active=True AND type='testing'")
+    ret = res.fetchone() != None
+    con.close()
+    return ret
+
+def add_premium_user(telegram_user_id, until):
     found = False
     con = sqlite3.connect(DB)
     cur = con.cursor()
     res = cur.execute(f"SELECT * FROM users WHERE telegram_user_id={telegram_user_id}")
     if res.fetchone() is None:
-        cur.execute(f"INSERT INTO users VALUES ({telegram_user_id}, True)")
+        cur.execute(f"INSERT INTO users VALUES ({telegram_user_id}, True, 'premium', '{until}')")
         con.commit()
     else:
-        cur.execute(f"UPDATE users SET active = True WHERE telegram_user_id={telegram_user_id}")
+        cur.execute(f"UPDATE users SET active = True, type = 'premium', until = {until} WHERE telegram_user_id={telegram_user_id}")
         con.commit()
         found = True
     con.close()
     return found
 
+def add_test_user(telegram_user_id, until):
+    found = False
+    con = sqlite3.connect(DB)
+    cur = con.cursor()
+    res = cur.execute(f"SELECT * FROM users WHERE telegram_user_id={telegram_user_id}")
+    if res.fetchone() is None:
+        cur.execute(f"INSERT INTO users VALUES ({telegram_user_id}, True, 'testing', '{until}')")
+        con.commit()
+    else:
+        found = True
+    con.close()
+    return not found
+
 def remove_valid_user(telegram_user_id):
     con = sqlite3.connect(DB)
     cur = con.cursor()
@@ -135,3 +164,11 @@ def remove_product(product):
         con.close()
         removed = True
     return removed
+
+def count_user_products(telegram_user_id):
+    con = sqlite3.connect(DB)
+    cur = con.cursor()
+    res = cur.execute(f"SELECT Count() FROM products WHERE telegram_user_id={telegram_user_id}")
+    ret = res.fetchone()[0]
+    con.close()
+    return ret

wallamanta/wallamanta.py

@@ -39,97 +39,115 @@ async def help_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> No
     Los campos opcionales que se dejen vacíos tomarán el valor configurado en el archivo `.env`\n
     Lista los productos con `/list` o obtén la información de uno en concreto con `/list nombre del producto`\n
     Borra un producto con `/remove nombre del producto`"""
-        await update.message.reply_markdown_v2(helpers.telegram_escape_characters(message))
+    else:
+        message = """Activa tu periodo de prueba de 7 días con `/test` o contacta con @jocarduck para más información."""    
+    await update.message.reply_markdown_v2(helpers.telegram_escape_characters(message))
 
 async def add_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
     telegram_user_id = helpers.get_telegram_user_id(update)
     if walladb.is_user_valid(telegram_user_id):
+        number_of_products = walladb.count_user_products(telegram_user_id)
         message = """Tienes que pasar el número correcto de parámetros: `/add producto;precio_mínimo;precio_máximo,excluir_título(opcional, separado por comas);excluir_descripción_y_título(opciona, separado por comas);latitud(opcional);longitud(opcional),distancia(opcional)`\n
     Ejemplo: `/add placa base itx;0;150`\n
     Ejemplo 2: `/add cpu;10;30;;intel,core 2 duo,celeron;;;100`\n
     Los campos opcionales que se dejen vacíos tomarán el valor configurado en el archivo `.env`"""
-
+        valid = False
-        args = update.message.text.split("/add ")
+        if walladb.is_user_testing(telegram_user_id):
-        if len(args) == 1:
+            valid = True
-            pass
+            if number_of_products > 5:
-        elif len(args[1].split(";")) > 2:
+                message = "Tienes más de 5 productos en seguimiento. Con premium puedes tener hasta 20."
-            product = dict()
+                valid = False
-            product['telegram_user_id'] = telegram_user_id
+        elif walladb.is_user_premium(telegram_user_id):
-            args = args[1].split(";")
+            valid = True
-            product['product_name'], product['min_price'], product['max_price'] = args[0:3]
+            if number_of_products > 20:
-            if len(args) > 3 and args[3]:
+                message = "Tienes más de 20 productos en seguimiento. Borra algunos para añadir más."
-                product['title_exclude'] = args[3]
+                valid = False
-            if len(args) > 4 and args[4]:
+        if valid:
-                product['title_description_exclude'] = args[4]
+            args = update.message.text.split("/add ")
-            if len(args) > 5 and args[5]:
+            if len(args) == 1:
-                product['latitude'] = args[5]
+                pass
-            if len(args) > 6 and args[6]:
+            elif len(args[1].split(";")) > 2:
-                product['longitude'] = args[6]
+                product = dict()
-            if len(args) > 7 and args[7]:
+                product['telegram_user_id'] = telegram_user_id
-                product['distance'] = args[7]
+                args = args[1].split(";")
-            
+                product['product_name'], product['min_price'], product['max_price'] = args[0:3]
-            logging.info(f'Adding: {product}')
+                if len(args) > 3 and args[3]:
-            if not walladb.get_product(product):
+                    product['title_exclude'] = args[3]
-                walladb.add_product(product)
+                if len(args) > 4 and args[4]:
-                p = threading.Thread(target=Worker.run, args=(walladb.get_product(product), ))
+                    product['title_description_exclude'] = args[4]
-                p.start()
+                if len(args) > 5 and args[5]:
-                message = f"Añadido {walladb.get_product(product)['product_name']} a seguimiento"
+                    product['latitude'] = args[5]
-            else:
+                if len(args) > 6 and args[6]:
-                message = f"{walladb.get_product(product)['product_name']} ya está en seguimiento!"
+                    product['longitude'] = args[6]
+                if len(args) > 7 and args[7]:
+                    product['distance'] = args[7]
+                
+                logging.info(f'Adding: {product}')
+                if not walladb.get_product(product):
+                    walladb.add_product(product)
+                    p = threading.Thread(target=Worker.run, args=(walladb.get_product(product), ))
+                    p.start()
+                    message = f"Añadido {walladb.get_product(product)['product_name']} a seguimiento."
+                else:
+                    message = f"¡{walladb.get_product(product)['product_name']} ya está en seguimiento!"
         await update.message.reply_markdown_v2(helpers.telegram_escape_characters(message))
 
 async def remove_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
     telegram_user_id = helpers.get_telegram_user_id(update)
     if walladb.is_user_valid(telegram_user_id):
-        product_to_remove = update.message.text[len('/remove '):]
+        if walladb.is_user_testing(telegram_user_id) or walladb.is_user_premium(telegram_user_id):
-        message = f"{product_to_remove} no está en seguimiento!"
+            product_to_remove = update.message.text[len('/remove '):]
-        if walladb.remove_product({'product_name' : product_to_remove, \
+            message = f"¡{product_to_remove} no está en seguimiento!"
-                                'telegram_user_id' : telegram_user_id}):
+            if walladb.remove_product({'product_name' : product_to_remove, \
-            message = f"{product_to_remove} borrado!"
+                                    'telegram_user_id' : telegram_user_id}):
-        await update.message.reply_text(message)
+                message = f"¡{product_to_remove} borrado de la lista de seguimiento!"
+            await update.message.reply_text(message)
 
 async def list_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
     telegram_user_id = helpers.get_telegram_user_id(update)
     if walladb.is_user_valid(telegram_user_id):
-        products = walladb.get_products_from_user(telegram_user_id)
+        if walladb.is_user_testing(telegram_user_id) or walladb.is_user_premium(telegram_user_id):
+            products = walladb.get_products_from_user(telegram_user_id)
 
-        args = update.message.text.split("/list ")
+            args = update.message.text.split("/list ")
-        found = False
+            found = False
-        if len(args) > 1:
+            if len(args) > 1:
-            product = walladb.get_product({'product_name':args[1],'telegram_user_id':telegram_user_id})
+                product = walladb.get_product({'product_name':args[1],'telegram_user_id':telegram_user_id})
-            if product:
+                if product:
-                table = prettytable.PrettyTable(['Campo', 'Valor'])
+                    table = prettytable.PrettyTable(['Campo', 'Valor'])
-                table.align['Campo'] = 'l'
+                    table.align['Campo'] = 'l'
-                table.align['Valor'] = 'r'
+                    table.align['Valor'] = 'r'
-                for key in product:
+                    for key in product:
-                    table.add_row([key, product[key]])
+                        table.add_row([key, product[key]])
-            found = True
+                found = True
-        if not found:
+            if not found:
-            table = prettytable.PrettyTable(['Producto', 'Mín', 'Máx'])
+                table = prettytable.PrettyTable(['Producto', 'Mín', 'Máx'])
-            table.align['Producto'] = 'l'
+                table.align['Producto'] = 'l'
-            table.align['Mín'] = 'r'
+                table.align['Mín'] = 'r'
-            table.align['Máx'] = 'r'
+                table.align['Máx'] = 'r'
-            for product in products:
+                for product in products:
-                table.add_row([helpers.telegram_escape_characters(product['product_name']), f"{helpers.telegram_escape_characters(product['min_price'])}€", f"{helpers.telegram_escape_characters(product['max_price'])}€"])
+                    table.add_row([helpers.telegram_escape_characters(product['product_name']), f"{helpers.telegram_escape_characters(product['min_price'])}€", f"{helpers.telegram_escape_characters(product['max_price'])}€"])
-        await update.message.reply_markdown_v2(f'```{(table)}```')
+            await update.message.reply_markdown_v2(f'```{(table)}```')
 
 async def admin_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
     if helpers.is_user_admin(update.message.chat_id):
         await update.message.reply_markdown_v2(helpers.telegram_escape_characters("¡Eres admin!"))
 
-async def add_user_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
+async def add_premium_user_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
     if helpers.is_user_admin(update.message.chat_id):
-        telegram_user_id = update.message.text.split('/add_user ')[1]
+        telegram_user_id = update.message.text.split('/add_premium_user ')[1]
-        if walladb.add_valid_user(telegram_user_id):
+        days = update.message.text.split('/add_premium_user ')[2]
+        until = helpers.get_date_ahead(days)
+        if walladb.add_premium_user(telegram_user_id, until):
             products = walladb.get_products_from_user(telegram_user_id)
 
             for product in products:
                 logging.info(product)
                 p = threading.Thread(target=Worker.run, args=(product, ))
                 p.start()
-            await update.message.reply_markdown_v2(helpers.telegram_escape_characters(f"{telegram_user_id} re-activado. Re-activando productos."))
+            await update.message.reply_markdown_v2(helpers.telegram_escape_characters(f"{telegram_user_id} re-activado hasta {until}. Re-activando productos."))
         else:
-            await update.message.reply_markdown_v2(helpers.telegram_escape_characters(f"{telegram_user_id} añadido."))
+            await update.message.reply_markdown_v2(helpers.telegram_escape_characters(f"{telegram_user_id} añadido hasta {until}."))
 
 async def remove_user_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
     if helpers.is_user_admin(update.message.chat_id):
@@ -137,6 +155,26 @@ async def remove_user_command(update: Update, context: ContextTypes.DEFAULT_TYPE
         walladb.remove_valid_user(telegram_user_id)
         await update.message.reply_markdown_v2(helpers.telegram_escape_characters(f"{telegram_user_id} desactivado."))
 
+async def status_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
+    # TODO: Obtener estado del usuario y fecha de caducidad de suscripción en el caso de haberla
+    message = ''    
+    await update.message.reply_markdown_v2(helpers.telegram_escape_characters(message))
+
+async def test_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
+    # TODO: Activar usuario en modo testing y responder con fecha de caducidad
+    telegram_user_id = helpers.get_telegram_user_id(update)
+    if not walladb.is_user_valid(telegram_user_id):
+        until = helpers.get_date_ahead(7)
+        walladb.add_test_user(telegram_user_id, until)
+        message = f"Periodo de prueba activado hasta el {until}."
+    else:
+        message = "Ya has utilizado el periodo de prueba."
+        if walladb.is_user_testing(telegram_user_id):
+            message = "Ya estás en el periodo de prueba."
+        elif walladb.is_user_premium(telegram_user_id):
+            message = "Ya eres premium. No puedes volver al periodo de prueba."
+    await update.message.reply_markdown_v2(helpers.telegram_escape_characters(message))
+
 def main()->None:
     walladb.setup_db()
     products = walladb.get_all_products()
@@ -156,8 +194,10 @@ def main()->None:
     application.add_handler(CommandHandler("remove", remove_command))
     application.add_handler(CommandHandler("list", list_command))
     application.add_handler(CommandHandler("admin", admin_command))
-    application.add_handler(CommandHandler("add_user", add_user_command))
+    application.add_handler(CommandHandler("add_premium_user", add_premium_user_command))
     application.add_handler(CommandHandler("remove_user", remove_user_command))
+    application.add_handler(CommandHandler("status", status_command))
+    application.add_handler(CommandHandler("test", test_command))
 
     # on non command i.e message - echo the message on Telegram
     #application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, echo))

wallamanta/worker.py

@@ -23,6 +23,16 @@ logger = logging.getLogger(__name__)
 
 class Worker:
 
+    def is_valid_request(product):
+        is_valid = True
+        if not walladb.get_product(product):
+            is_valid = False
+        if not walladb.is_user_valid(product['telegram_user_id']):
+            is_valid = False
+        if not walladb.is_user_premium(product['telegram_user_id']):
+            is_valid = False
+        return is_valid
+
     def request(self, product_name, n_articles, latitude=LATITUDE, longitude=LONGITUDE, distance='0', condition='all', min_price=0, max_price=10000000):
         url = (f"http://api.wallapop.com/api/v3/general/search?keywords={product_name}"
                     f"&order_by=newest&latitude={latitude}"
@@ -53,6 +63,8 @@ class Worker:
 
     def first_run(self, product):
         list = []
+        if not self.is_valid_request(product):
+            return list
         articles = self.request(product['product_name'], 0, product['latitude'], product['longitude'], product['distance'], product['condition'], product['min_price'], product['max_price'])
         for article in articles: 
             list.insert(0, article['id'])
@@ -62,10 +74,8 @@ class Worker:
         exec_times = []
         bot = telegram.Bot(token = TELEGRAM_TOKEN)
         while True:
-            if not walladb.get_product(product):
+            if not self.is_valid_request(product):
-                break # Exits worker if product not in DB anymore
+                break # Exits and ends worker thread
-            if not walladb.is_user_valid(product['telegram_user_id']):
-                break # Exits worker if user is not active anymore
             start_time = time.time()
             articles = self.request(product['product_name'], 0, product['latitude'], product['longitude'], product['distance'], product['condition'], product['min_price'], product['max_price'])
             for article in articles: