diff --git a/bot.py b/bot.py
index 0a6d057..a86cbe9 100644
--- a/bot.py
+++ b/bot.py
@@ -33,6 +33,7 @@ def main_menu_kb(tg_id: int) -> InlineKeyboardMarkup:
rows = [
[InlineKeyboardButton(text="Получить VPN", callback_data="menu:get")],
[InlineKeyboardButton(text="Моя подписка", callback_data="menu:mysub")],
+ [InlineKeyboardButton(text="Пригласить друга", callback_data="menu:referral")],
[InlineKeyboardButton(text="О сервисе", callback_data="menu:about")],
]
if is_admin(tg_id):
@@ -40,6 +41,14 @@ def main_menu_kb(tg_id: int) -> InlineKeyboardMarkup:
return InlineKeyboardMarkup(inline_keyboard=rows)
+async def get_bot_username() -> str:
+ global _bot_username
+ if _bot_username is None:
+ me = await bot.get_me()
+ _bot_username = me.username
+ return _bot_username
+
+
def nodes_kb(prefix: str) -> InlineKeyboardMarkup:
rows = []
for n in db.list_nodes(enabled_only=True):
@@ -90,6 +99,12 @@ async def send_main_menu(message: Message):
async def start_deeplink(message: Message, command: CommandObject):
user = db.get_or_create_user(message.from_user.id, message.from_user.username)
payload = command.args or ""
+ if payload.startswith("ref_") or payload.startswith("ref-"):
+ ref_code = payload[4:]
+ referrer = db.get_user_by_ref_code(ref_code)
+ if referrer and settings.get_referral_settings()["enabled"]:
+ db.set_referred_by(message.from_user.id, referrer["tg_id"])
+ return await send_main_menu(message)
if payload.startswith("gift_") or payload.startswith("gift-"):
code = payload[5:]
gift, err = db.redeem_gift_code(code, message.from_user.id)
@@ -140,6 +155,32 @@ async def cb_about(cb: CallbackQuery):
await cb.answer()
+@dp.callback_query(F.data == "menu:referral")
+async def cb_referral(cb: CallbackQuery):
+ kb = InlineKeyboardMarkup(inline_keyboard=[[InlineKeyboardButton(text="Назад", callback_data="menu:main")]])
+ ref_settings = settings.get_referral_settings()
+ if not ref_settings["enabled"]:
+ await cb.message.edit_text("Реферальная программа сейчас отключена.", reply_markup=kb)
+ return await cb.answer()
+ user = db.get_or_create_user(cb.from_user.id, cb.from_user.username)
+ stats = db.referral_stats(cb.from_user.id)
+ username = await get_bot_username()
+ link = f"https://t.me/{username}?start=ref_{user['ref_code']}"
+ days = ref_settings["bonus_days"]
+ text = (
+ f"Пригласи друга\n\n"
+ f"За каждого друга, который активирует подписку по твоей ссылке, "
+ f"вы оба получаете +{days} дн. к подписке.\n\n"
+ f"{DIVIDER}\n"
+ f"Твоя ссылка:\n{link}\n\n"
+ f"Приглашено: {stats['referred_count']}\n"
+ )
+ if stats["bonus_days_pending"]:
+ text += f"Накоплено бонусных дней (зачислятся при следующей подписке): {stats['bonus_days_pending']}\n"
+ await cb.message.edit_text(text, reply_markup=kb)
+ await cb.answer()
+
+
@dp.callback_query(F.data == "menu:get")
async def cb_get(cb: CallbackQuery):
await cb.message.edit_text("Выбери сервер:", reply_markup=nodes_kb("node"))
diff --git a/config.py b/config.py
index 7f67c5c..f5a1cf1 100644
--- a/config.py
+++ b/config.py
@@ -120,3 +120,6 @@ PLATEGA_SECRET = env("PLATEGA_SECRET", "")
HWID_LIMIT_ENABLED = env("HWID_LIMIT_ENABLED", "false").lower() == "true"
HWID_FALLBACK_LIMIT = int(env("HWID_FALLBACK_LIMIT", "3"))
+
+REFERRAL_ENABLED = env("REFERRAL_ENABLED", "true").lower() == "true"
+REFERRAL_BONUS_DAYS = int(env("REFERRAL_BONUS_DAYS", "3"))
diff --git a/db.py b/db.py
index 68dd516..62dec11 100644
--- a/db.py
+++ b/db.py
@@ -133,6 +133,10 @@ _NEW_NODE_COLUMNS = {
_NEW_USER_COLUMNS = {
"hwid_limit": "INTEGER",
+ "ref_code": "TEXT",
+ "referred_by": "INTEGER",
+ "referral_rewarded": "INTEGER NOT NULL DEFAULT 0",
+ "bonus_days_pending": "INTEGER NOT NULL DEFAULT 0",
}
_NEW_ADMIN_SESSION_COLUMNS = {
@@ -159,6 +163,7 @@ def _migrate():
for name, decl in _NEW_USER_COLUMNS.items():
if name not in ucols:
conn.execute(f"ALTER TABLE users ADD COLUMN {name} {decl}")
+ conn.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_users_ref_code ON users(ref_code)")
scols = {r["name"] for r in conn.execute("PRAGMA table_info(admin_sessions)").fetchall()}
for name, decl in _NEW_ADMIN_SESSION_COLUMNS.items():
if name not in scols:
@@ -353,22 +358,97 @@ def delete_node(code: str):
conn.execute("DELETE FROM nodes WHERE code=?", (code,))
+def _generate_ref_code(conn) -> str:
+ for _ in range(20):
+ code = secrets.token_hex(4)
+ if not conn.execute("SELECT 1 FROM users WHERE ref_code=?", (code,)).fetchone():
+ return code
+ raise RuntimeError("could not generate a unique ref_code")
+
+
def get_or_create_user(tg_id: int, username: str | None):
with get_conn() as conn:
row = conn.execute("SELECT * FROM users WHERE tg_id=?", (tg_id,)).fetchone()
if row:
if username and row["username"] != username:
conn.execute("UPDATE users SET username=? WHERE tg_id=?", (username, tg_id))
+ if not row["ref_code"]:
+ conn.execute(
+ "UPDATE users SET ref_code=? WHERE tg_id=?", (_generate_ref_code(conn), tg_id)
+ )
+ row = conn.execute("SELECT * FROM users WHERE tg_id=?", (tg_id,)).fetchone()
return dict(row)
token = secrets.token_hex(16)
+ ref_code = _generate_ref_code(conn)
conn.execute(
- "INSERT INTO users (tg_id, token, username, created_at) VALUES (?,?,?,?)",
- (tg_id, token, username, now_iso()),
+ "INSERT INTO users (tg_id, token, username, ref_code, created_at) VALUES (?,?,?,?,?)",
+ (tg_id, token, username, ref_code, now_iso()),
)
row = conn.execute("SELECT * FROM users WHERE tg_id=?", (tg_id,)).fetchone()
return dict(row)
+def get_user_by_ref_code(ref_code: str):
+ with get_conn() as conn:
+ row = conn.execute("SELECT * FROM users WHERE ref_code=?", (ref_code,)).fetchone()
+ return dict(row) if row else None
+
+
+def set_referred_by(tg_id: int, referrer_tg_id: int) -> bool:
+ """First-touch attribution: only takes effect for a brand-new account
+ (no subscriptions yet) that isn't already attributed, and never to self."""
+ if tg_id == referrer_tg_id:
+ return False
+ with get_conn() as conn:
+ row = conn.execute("SELECT referred_by FROM users WHERE tg_id=?", (tg_id,)).fetchone()
+ if not row or row["referred_by"] is not None:
+ return False
+ has_sub = conn.execute("SELECT 1 FROM subscriptions WHERE tg_id=?", (tg_id,)).fetchone()
+ if has_sub:
+ return False
+ referrer = conn.execute("SELECT 1 FROM users WHERE tg_id=?", (referrer_tg_id,)).fetchone()
+ if not referrer:
+ return False
+ conn.execute("UPDATE users SET referred_by=? WHERE tg_id=?", (referrer_tg_id, tg_id))
+ return True
+
+
+def _apply_bonus_days(conn, tg_id: int, days: int):
+ if days <= 0:
+ return
+ row = conn.execute(
+ "SELECT uuid, expires_at FROM subscriptions WHERE tg_id=? AND active=1 AND held_at IS NULL "
+ "ORDER BY expires_at DESC LIMIT 1",
+ (tg_id,),
+ ).fetchone()
+ if row:
+ new_expires = datetime.datetime.fromisoformat(row["expires_at"]) + datetime.timedelta(days=days)
+ conn.execute("UPDATE subscriptions SET expires_at=? WHERE uuid=?", (new_expires.isoformat(), row["uuid"]))
+ else:
+ conn.execute(
+ "UPDATE users SET bonus_days_pending = COALESCE(bonus_days_pending, 0) + ? WHERE tg_id=?",
+ (days, tg_id),
+ )
+
+
+def credit_bonus_days(tg_id: int, days: int):
+ with get_conn() as conn:
+ _apply_bonus_days(conn, tg_id, days)
+
+
+def referral_stats(tg_id: int) -> dict:
+ with get_conn() as conn:
+ user = conn.execute("SELECT ref_code, bonus_days_pending FROM users WHERE tg_id=?", (tg_id,)).fetchone()
+ count = conn.execute(
+ "SELECT COUNT(*) c FROM users WHERE referred_by=? AND referral_rewarded=1", (tg_id,)
+ ).fetchone()["c"]
+ return {
+ "ref_code": user["ref_code"] if user else None,
+ "bonus_days_pending": user["bonus_days_pending"] if user else 0,
+ "referred_count": count,
+ }
+
+
def get_user_by_token(token: str):
with get_conn() as conn:
row = conn.execute("SELECT * FROM users WHERE token=?", (token,)).fetchone()
@@ -377,17 +457,32 @@ def get_user_by_token(token: str):
def create_subscription(tg_id: int, node: str, plan_days: int, plan_code: str, source: str = "bot", client_uuid: str | None = None):
import uuid as uuidlib
+ import config
cid = client_uuid or str(uuidlib.uuid4())
created = datetime.datetime.utcnow()
expires = created + datetime.timedelta(days=plan_days)
with get_conn() as conn:
+ is_first = conn.execute("SELECT 1 FROM subscriptions WHERE tg_id=?", (tg_id,)).fetchone() is None
+ urow = conn.execute(
+ "SELECT referred_by, referral_rewarded, bonus_days_pending FROM users WHERE tg_id=?", (tg_id,)
+ ).fetchone()
+ pending = urow["bonus_days_pending"] if urow else 0
+ if pending:
+ expires += datetime.timedelta(days=pending)
+ conn.execute("UPDATE users SET bonus_days_pending=0 WHERE tg_id=?", (tg_id,))
conn.execute(
"INSERT INTO subscriptions (uuid, tg_id, node, plan, created_at, expires_at, active, source) "
"VALUES (?,?,?,?,?,?,1,?)",
(cid, tg_id, node, plan_code, created.isoformat(), expires.isoformat(), source),
)
- return {"uuid": cid, "tg_id": tg_id, "node": node, "plan": plan_code, "expires_at": expires.isoformat()}
+ if is_first and urow and urow["referred_by"] and not urow["referral_rewarded"] and config.REFERRAL_ENABLED:
+ conn.execute("UPDATE users SET referral_rewarded=1 WHERE tg_id=?", (tg_id,))
+ bonus = config.REFERRAL_BONUS_DAYS
+ _apply_bonus_days(conn, tg_id, bonus)
+ _apply_bonus_days(conn, urow["referred_by"], bonus)
+ expires_final = conn.execute("SELECT expires_at FROM subscriptions WHERE uuid=?", (cid,)).fetchone()["expires_at"]
+ return {"uuid": cid, "tg_id": tg_id, "node": node, "plan": plan_code, "expires_at": expires_final}
def list_active_subscriptions(tg_id: int | None = None, node: str | None = None):
diff --git a/settings.py b/settings.py
index 3679ad5..a7394d5 100644
--- a/settings.py
+++ b/settings.py
@@ -87,3 +87,17 @@ def get_hwid_settings() -> dict:
"enabled": _bool(raw.get("HWID_LIMIT_ENABLED"), config.HWID_LIMIT_ENABLED),
"fallback_limit": limit if limit > 0 else config.HWID_FALLBACK_LIMIT,
}
+
+
+def get_referral_settings() -> dict:
+ raw = legal.read_env_vars(["REFERRAL_ENABLED", "REFERRAL_BONUS_DAYS"])
+ days = _positive_int(raw.get("REFERRAL_BONUS_DAYS"), config.REFERRAL_BONUS_DAYS)
+ return {
+ "enabled": _bool(raw.get("REFERRAL_ENABLED"), config.REFERRAL_ENABLED),
+ "bonus_days": days if days > 0 else config.REFERRAL_BONUS_DAYS,
+ }
+
+
+def set_referral_settings(enabled: bool, bonus_days: int):
+ legal.update_env_var("REFERRAL_ENABLED", "true" if enabled else "false")
+ legal.update_env_var("REFERRAL_BONUS_DAYS", str(int(bonus_days)))