Empezando

2023-03-05 00:01:43 +01:00
parent 2abff6dddb
commit 90ea99ca3a
3 changed files with 157 additions and 103 deletions
--- a/wallamanta/alert.py
+++ b/wallamanta/alert.py
@@ -32,100 +32,89 @@ def save_json_file(products):
        json.dump(products, outfile, indent=2)

 async def help_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
-    """Send a message when the command /help is issued."""
-    message = """Añade un producto con `/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`\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))
+    if walladb.is_user_valid(helpers.get_telegram_user_id(update)):
+        message = """Añade un producto con `/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`\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))

 async def add_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
-    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`"""
+    if walladb.is_user_valid(helpers.get_telegram_user_id(update)):
+        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`"""

-    args = update.message.text.split("/add ")
-    if len(args) == 1:
-       pass
-    elif len(args[1].split(";")) > 2:
-        args = args[1].split(";")
-        logging.info(f'Adding: {args}')
-        title_exclude, title_description_exclude, latitude, longitude, distance = [], [], LATITUDE, LONGITUDE, "0"
-        product_name, min_price, max_price = args[0:3]
-        if len(args) > 3 and args[3]:
-            title_exclude = args[3].split(",")
-        if len(args) > 4 and args[4]:
-            title_description_exclude = args[4].split(",")
-        if len(args) > 5 and args[5]:
-            latitude = args[5]
-        if len(args) > 6 and args[6]:
-            longitude = args[6]
-        if len(args) > 7 and args[7]:
-            distance = args[7]
-
-        product = {"product_name": product_name,
-                    "distance": distance,
-                    "latitude": latitude,
-                    "longitude": longitude,
-                    "condition": "all",
-                    "min_price": min_price,
-                    "max_price": max_price,
-                    "title_keyword_exclude" : title_exclude,
-                    "exclude": title_description_exclude
-                    }
-        
-        products = parse_json_file()
-        found = False
-        for fproduct in products:
-            if fproduct['product_name'] == product_name:
-                found = True
-        if not found:
-            products.append(product)
-            save_json_file(products)
-            p = threading.Thread(target=Worker.run, args=(product, ))
-            p.start()
-            message = f"Añadido {product_name} a seguimiento"
-        else:
-            message = f"{product_name} ya está en seguimiento!"
-    await update.message.reply_markdown_v2(helpers.telegram_escape_characters(message))
+        args = update.message.text.split("/add ")
+        if len(args) == 1:
+            pass
+        elif len(args[1].split(";")) > 2:
+            args = args[1].split(";")
+            logging.info(f'Adding: {args}')
+            if not walladb.get_product(args, helpers.get_telegram_user_id(update)):
+                walladb.add_product(args, helpers.get_telegram_user_id(update))
+                p = threading.Thread(target=Worker.run, args=(walladb.get_product(args, helpers.get_telegram_user_id(update)), ))
+                p.start()
+                message = f"Añadido {walladb.get_product(args, helpers.get_telegram_user_id(update))['product_name']} a seguimiento"
+            else:
+                message = f"{walladb.get_product(args, helpers.get_telegram_user_id(update))['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:
-    product_to_remove = update.message.text[len('/remove '):]
-    message = f"{product_to_remove} no está en seguimiento!"
-    products = parse_json_file()
-    for product in products:
-        if product['product_name'] == product_to_remove:
-            products.remove(product)
-            message = f"{product_to_remove} borrado!"
-    save_json_file(products)
-    await update.message.reply_text(message)
+    if walladb.is_user_valid(helpers.get_telegram_user_id(update)):
+        product_to_remove = update.message.text[len('/remove '):]
+        message = f"{product_to_remove} no está en seguimiento!"
+        products = parse_json_file()
+        for product in products:
+            if product['product_name'] == product_to_remove:
+                products.remove(product)
+                message = f"{product_to_remove} borrado!"
+        save_json_file(products)
+        await update.message.reply_text(message)

 async def list_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
-    products = parse_json_file()
+    if walladb.is_user_valid(helpers.get_telegram_user_id(update)):
+        products = parse_json_file()

-    args = update.message.text.split("/list ")
-    found = False
-    if len(args) > 1:
-        table = prettytable.PrettyTable(['Campo', 'Valor'])
-        table.align['Campo'] = 'l'
-        table.align['Valor'] = 'r'
-        for product in products:
-            if product['product_name'] == args[1]:
-                for key in product:
-                    table.add_row([key, product[key]])
-                found = True
-                break
-    if not found:
-        table = prettytable.PrettyTable(['Producto', 'Mín', 'Máx'])
-        table.align['Producto'] = 'l'
-        table.align['Mín'] = 'r'
-        table.align['Máx'] = 'r'
-        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'])}€"])
-    await update.message.reply_markdown_v2(f'```{(table)}```')
+        args = update.message.text.split("/list ")
+        found = False
+        if len(args) > 1:
+            table = prettytable.PrettyTable(['Campo', 'Valor'])
+            table.align['Campo'] = 'l'
+            table.align['Valor'] = 'r'
+            for product in products:
+                if product['product_name'] == args[1]:
+                    for key in product:
+                        table.add_row([key, product[key]])
+                    found = True
+                    break
+        if not found:
+            table = prettytable.PrettyTable(['Producto', 'Mín', 'Máx'])
+            table.align['Producto'] = 'l'
+            table.align['Mín'] = 'r'
+            table.align['Máx'] = 'r'
+            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'])}€"])
+        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:
+    if helpers.is_user_admin(update.message.chat_id):
+        telegram_user_id = update.message.text.split('/add_user ')[1]
+        walladb.add_valid_user(telegram_user_id)
+        await update.message.reply_markdown_v2(helpers.telegram_escape_characters(f"{telegram_user_id} añadido."))
+
+async def remove_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('/remove_user ')[1]
+        walladb.remove_valid_user(telegram_user_id)
+        await update.message.reply_markdown_v2(helpers.telegram_escape_characters(f"{telegram_user_id} desactivado."))

 def main()->None:
    walladb.setup_db()
@@ -145,6 +134,7 @@ def main()->None:
    application.add_handler(CommandHandler("add", add_command))
    application.add_handler(CommandHandler("remove", remove_command))
    application.add_handler(CommandHandler("list", list_command))
+    application.add_handler(CommandHandler("admin", list_command))

    # on non command i.e message - echo the message on Telegram
    #application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, echo))
--- a/wallamanta/helpers.py
+++ b/wallamanta/helpers.py
@@ -1,5 +1,6 @@
 #TELEGRAM_ESCAPE_CHARACTERS = ['_', '*', '[', ']', '(', ')', '~', '>', '#', '+', '-', '=', '|', '{', '}', '.', '!']
 TELEGRAM_ESCAPE_CHARACTERS = ['_', '*', '[', ']', '(', ')', '~', '>', '+', '-', '=', '|', '{', '}', '.', '!']
+ADMIN_IDS = [10101691]

 def telegram_escape_characters(text):
    for character in TELEGRAM_ESCAPE_CHARACTERS:
@@ -7,4 +8,10 @@ def telegram_escape_characters(text):
            text = text.replace(character, f'\\{character}')
        except:
            pass
-    return text
+    return text
+
+def is_user_admin(telegram_user_id):
+    return telegram_user_id in ADMIN_IDS
+
+def get_telegram_user_id(update):
+    return update.message.chat_id
--- a/wallamanta/walladb.py
+++ b/wallamanta/walladb.py
@@ -1,16 +1,73 @@
-import sqlite3
-
-DB = "/app/data/wallamanta.db"
-
-def setup_db():
-    con = sqlite3.connect(DB)
-    cur = con.cursor()
-    cur.execute("CREATE TABLE IF NOT EXISTS users(telegram_user_id, telegram_user, active)")
-    cur.execute("CREATE TABLE IF NOT EXISTS products(product_name, distance, latitude, longitude, condition, min_price, max_price, title_keyword_exclude, exclude, telegram_user_id)")
-    con.close()
-
-def is_user_valid(telegram_user_id):
-    con = sqlite3.connect(DB)
-    cur = con.cursor()
-    res = cur.execute(f"SELECT * FROM users WHERE telegram_user_id={telegram_user_id}")
-    return res.fetchone() is None
+import sqlite3
+
+DB = "/app/data/wallamanta.db"
+
+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 products(product_name, distance, latitude, longitude, condition, min_price, max_price, title_keyword_exclude, exclude, telegram_user_id)")
+    con.close()
+
+def is_user_valid(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")
+    ret = res.fetchone() is None
+    con.close()
+    return ret
+
+def add_valid_user(telegram_user_id):
+    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")
+        con.commit()
+    else:
+        cur.execute(f"UPDATE users SET active = True WHERE telegram_user_id={telegram_user_id}")
+        con.commit()
+    con.close()
+
+def remove_valid_user(telegram_user_id):
+    con = sqlite3.connect(DB)
+    cur = con.cursor()
+    res = cur.execute(f"SELECT * FROM users WHERE telegram_user_id={telegram_user_id}")
+    if res.fetchone() != None:
+        cur.execute(f"UPDATE users SET active = False WHERE telegram_user_id={telegram_user_id}")
+        con.commit()
+    con.close()
+
+def get_user_list():
+    con = sqlite3.connect(DB)
+    cur = con.cursor()
+    res = cur.execute(f"SELECT * FROM users")
+    ret = res.fetchall()
+    con.close()
+    return ret
+
+def get_product(args, telegram_user_id):
+    product_name = args[0]
+    con = sqlite3.connect(DB)
+    con.row_factory = sqlite3.Row
+    cur = con.cursor()
+    res = cur.execute(f"SELECT * FROM products WHERE telegram_user_id={telegram_user_id} AND product_name={product_name}")
+    ret = res.fetchone()
+    con.close()
+    return ret
+
+def add_product(args, telegram_user_id):
+    title_exclude, title_description_exclude, latitude, longitude, distance = [], [], LATITUDE, LONGITUDE, "0"
+    product_name, min_price, max_price = args[0:3]
+    if len(args) > 3 and args[3]:
+        title_exclude = args[3].split(",")
+    if len(args) > 4 and args[4]:
+        title_description_exclude = args[4].split(",")
+    if len(args) > 5 and args[5]:
+        latitude = args[5]
+    if len(args) > 6 and args[6]:
+        longitude = args[6]
+    if len(args) > 7 and args[7]:
+        distance = args[7]
+
+