diff --git a/src/background_tasks.py b/src/background_tasks.py deleted file mode 100644 index 295c480..0000000 --- a/src/background_tasks.py +++ /dev/null @@ -1,40 +0,0 @@ -import asyncio -import time -import logging -from .rate_limiter import cleanup_old_records - -logger = logging.getLogger(__name__) - -async def periodic_cleanup(): - while True: - try: - await asyncio.sleep(3600) - cleanup_old_records() - logger.info(f"Periodic cleanup completed at {time.strftime('%Y-%m-%d %H:%M:%S')}") - except asyncio.CancelledError: - logger.info("Periodic cleanup task cancelled") - break - except Exception as e: - logger.error(f"Error in periodic cleanup: {e}") - -async def startup_tasks(): - try: - cleanup_old_records() - logger.info("Initial cleanup completed") - cleanup_task = asyncio.create_task(periodic_cleanup()) - return cleanup_task - except Exception as e: - logger.error(f"Error in startup tasks: {e}") - return None - -async def shutdown_tasks(cleanup_task): - try: - if cleanup_task and not cleanup_task.done(): - cleanup_task.cancel() - try: - await cleanup_task - except asyncio.CancelledError: - pass - logger.info("Background tasks shutdown completed") - except Exception as e: - logger.error(f"Error in shutdown tasks: {e}") diff --git a/src/main.py b/src/main.py index 6c6cb1a..b34136a 100644 --- a/src/main.py +++ b/src/main.py @@ -1,4 +1,4 @@ -from fastapi import FastAPI, Request, HTTPException +from fastapi import FastAPI, Request, HTTPException, Response from fastapi.responses import RedirectResponse, HTMLResponse from fastapi.staticfiles import StaticFiles from fastapi.middleware.cors import CORSMiddleware @@ -10,31 +10,19 @@ from .models import initialize_database, close_database from .schemas import BadgeParams, UrlStatsResponse, SystemStatsResponse from .services import update_visit_count, get_url_visit_count, get_system_statistics, get_app_info, load_template -from .utils import get_client_ip, get_ip_hash, build_shields_url, get_security_headers -from .rate_limiter import get_rate_limit_info -from .background_tasks import startup_tasks, shutdown_tasks +from .utils import build_shields_url, get_security_headers LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO").upper() logging.basicConfig(level=LOG_LEVEL, format='%(asctime)s - %(levelname)s - %(name)s - %(message)s') logger = logging.getLogger(__name__) -cleanup_task = None - @asynccontextmanager async def lifespan(app: FastAPI): - global cleanup_task - logger.info("Starting BadgeTrack application...") - if not initialize_database(): raise RuntimeError("Failed to initialize database") - - cleanup_task = await startup_tasks() - yield - logger.info("Shutting down BadgeTrack application...") - await shutdown_tasks(cleanup_task) close_database() def get_app_version(): @@ -82,6 +70,7 @@ def get_app_version(): @app.get("/badge") async def badge( request: Request, + response: Response, url: str, label: str = "visits", color: str = "4ade80", @@ -93,11 +82,12 @@ async def badge( except Exception: raise HTTPException(status_code=400, detail="Invalid parameters.") - ip = get_client_ip(request) - ip_hash = get_ip_hash(ip) + cookie_id = request.cookies.get("visitor_id") try: - count, was_incremented = update_visit_count(ip_hash, params.url) + count, was_incremented, new_cookie_id = update_visit_count(cookie_id, params.url) + if new_cookie_id: + response.set_cookie(key="visitor_id", value=new_cookie_id, max_age=31536000, httponly=True, samesite="Lax") except ValueError as e: raise HTTPException(status_code=429, detail=str(e)) except Exception as e: @@ -129,13 +119,12 @@ async def get_url_stats_endpoint(url: str): async def get_system_stats_endpoint(): try: stats = get_system_statistics() - rate_info = get_rate_limit_info() return SystemStatsResponse( total_tracked_urls=stats["total_tracked_urls"], total_visits=stats["total_visits"], new_badges_today=stats["new_badges_today"], - rate_limit_window_hours=rate_info["rate_limit_window_hours"] + rate_limit_window_hours=0 ) except Exception as e: logger.error(f"Error getting system stats: {e}") diff --git a/src/models.py b/src/models.py index 45a6bac..05c3206 100644 --- a/src/models.py +++ b/src/models.py @@ -17,26 +17,22 @@ class BaseModel(Model): class Meta: database = db -class IpHash(BaseModel): - """Store unique IP hashes""" - hash = CharField(max_length=64, unique=True) - class Badge(BaseModel): """Badge URLs with their total visit counts""" url = CharField(max_length=200, unique=True) visits = IntegerField(default=0) created = IntegerField() # When first created -class Visit(BaseModel): - """Track individual IP visits to prevent spam""" - ip_hash = ForeignKeyField(IpHash, backref='visits') - badge = ForeignKeyField(Badge, backref='visits') +class Cookie(BaseModel): + """Track individual cookie visits to prevent spam""" + cookie_id = CharField(max_length=64, unique=True) + badge = ForeignKeyField(Badge, backref='cookies') last_visit = IntegerField() - + class Meta: database = db indexes = ( - (("ip_hash", "badge"), True), # One record per IP+Badge combo + (("cookie_id", "badge"), True), # One record per cookie+Badge combo ) def initialize_database(): @@ -49,7 +45,7 @@ def initialize_database(): if db.is_closed(): db.connect() - db.create_tables([IpHash, Badge, Visit], safe=True) + db.create_tables([Badge, Cookie], safe=True) logger.info("Database initialized successfully") return True except Exception as e: diff --git a/src/rate_limiter.py b/src/rate_limiter.py deleted file mode 100644 index 2437e01..0000000 --- a/src/rate_limiter.py +++ /dev/null @@ -1,59 +0,0 @@ -from .models import db, IpHash, Badge, Visit -import time -import logging -import os - -logger = logging.getLogger(__name__) - -RATE_LIMIT_WINDOW_SECONDS = int(os.getenv("RATE_LIMIT_WINDOW_SECONDS", "172800")) - -def is_visit_allowed(ip_hash_str: str, url_str: str) -> bool: - """Check if a visit from this IP to this URL is allowed based on rate limiting""" - now = int(time.time()) - try: - # Get IP hash record - ip_hash = IpHash.get_or_none(IpHash.hash == ip_hash_str) - if not ip_hash: - return True # First time seeing this IP, allow visit - - # Get badge record - badge = Badge.get_or_none(Badge.url == url_str) - if not badge: - return True # First time seeing this URL, allow visit - - # Check if there's a recent visit from this IP to this badge - visit_record = Visit.get_or_none( - (Visit.ip_hash == ip_hash) & (Visit.badge == badge) - ) - - if not visit_record: - return True # No previous visit from this IP to this badge - - # Check if enough time has passed since last visit - return (now - visit_record.last_visit) >= RATE_LIMIT_WINDOW_SECONDS - except Exception as e: - logger.error(f"Error checking visit rate limit: {e}") - return True # Allow visit on error to avoid blocking legitimate users - -def cleanup_old_records(): - """Clean up old records to save space""" - now = int(time.time()) - week_ago = now - 604800 # 7 days - - try: - with db.atomic(): - # Clean up old visit records - deleted_visits = Visit.delete().where( - Visit.last_visit < week_ago - ).execute() - - logger.info(f"Cleaned up {deleted_visits} old visit records.") - - except Exception as e: - logger.error(f"Error during cleanup: {e}") - -def get_rate_limit_info() -> dict: - return { - "rate_limit_window_hours": RATE_LIMIT_WINDOW_SECONDS // 3600, - "cleanup_retention_days": 7 - } diff --git a/src/services.py b/src/services.py index 6c26a49..63e570e 100644 --- a/src/services.py +++ b/src/services.py @@ -1,59 +1,50 @@ from typing import Tuple from peewee import fn -from .models import db, IpHash, Badge, Visit -from .rate_limiter import is_visit_allowed +from .models import db, Badge, Cookie import time import os import json import logging +import secrets logger = logging.getLogger(__name__) -def update_visit_count(ip_hash_str: str, url_str: str) -> Tuple[int, bool]: - """Update visit count for a URL and IP combination""" +def update_visit_count(cookie_id: str, url_str: str) -> Tuple[int, bool, str]: + """Update visit count for a URL and cookie combination""" current_time = int(time.time()) - + new_cookie_id = None + try: with db.atomic(): - # Get or create IP hash record - ip_hash, _ = IpHash.get_or_create(hash=ip_hash_str) - - # Get or create badge record badge, badge_created = Badge.get_or_create( url=url_str, defaults={'created': current_time} ) - - # Check if visit is allowed (rate limiting) - if is_visit_allowed(ip_hash_str, url_str): - # Get or create visit record for this IP-Badge combination - visit, visit_created = Visit.get_or_create( - ip_hash=ip_hash, - badge=badge, - defaults={'last_visit': current_time} - ) - - # Only increment if this is a new visit or enough time has passed - if visit_created or (current_time - visit.last_visit >= int(os.getenv("RATE_LIMIT_WINDOW_SECONDS", "172800"))): - visit.last_visit = current_time - visit.save() - - badge.visits += 1 - badge.save() - - return badge.visits, True - - # Visit not allowed or not enough time passed, return current count - return badge.visits, False - + + if not cookie_id: + cookie_id = secrets.token_hex(16) + new_cookie_id = cookie_id + + cookie, cookie_created = Cookie.get_or_create( + cookie_id=cookie_id, + badge=badge, + defaults={'last_visit': current_time} + ) + + if cookie_created: + badge.visits += 1 + badge.save() + return badge.visits, True, new_cookie_id + else: + return badge.visits, False, new_cookie_id + except Exception as e: logger.error(f"Error updating visit count: {e}") - # Try to get current count even if update failed try: badge = Badge.get(Badge.url == url_str) - return badge.visits, False + return badge.visits, False, new_cookie_id except Badge.DoesNotExist: - return 0, False + return 0, False, new_cookie_id def get_url_visit_count(url_str: str) -> int: """Get total visit count for a URL""" diff --git a/src/utils.py b/src/utils.py index 7d428da..2f40a60 100644 --- a/src/utils.py +++ b/src/utils.py @@ -18,19 +18,6 @@ logging.warning("Provided SECRET_KEY is too short. Consider generating a new one.") -def get_ip_hash(ip: str) -> str: - return hashlib.sha256((ip + SECRET_KEY).encode()).hexdigest() - - -def get_client_ip(request) -> str: - ip = ( - request.headers.get("x-forwarded-for", request.client.host) - .split(",")[0] - .strip() - ) - return ip - - def build_shields_url(label: str, count: int, color: str, style: str, logo: str = "") -> str: shields_url = ( f"https://img.shields.io/badge/" diff --git a/templates/about.html b/templates/about.html index 68108cf..21304f0 100644 --- a/templates/about.html +++ b/templates/about.html @@ -41,31 +41,19 @@

What is BadgeTrack?

dynamic visit counter badges for your projects. It provides real-time visit tracking with beautiful, customizable badges that can be embedded in your README files, websites, and documentation. -

How It Works

+

+

How It Works

The application uses FastAPI and SQLite to:

    -
  1. Track visits to your specified URLs
    2. +
  2. Track visits to your specified URLs using anonymous cookies
  3. Generate dynamic badges via Shields.io integration
  4. Provide real-time visit count updates
  5. Offer customizable styles, colors, and logos
    6. -
  6. Implement enhanced rate limiting for abuse prevention
-

Rate Limiting

-

- BadgeTrack implements a dual rate limiting system: -

- -

Privacy & Security

- User privacy is prioritized. IP addresses are hashed using a secure - key and are never stored in plain text. The application implements - rate limiting to prevent abuse while ensuring accurate visit - tracking. No personal data is collected or stored. + User privacy is prioritized. The application uses an anonymous, unique cookie to distinguish between visits and prevent a single user from incrementing the counter multiple times for the same badge. No personal data is collected or stored. The cookie contains no identifying information.

Open Source