Refactor visit tracking system to use cookies instead of IP addresses.

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}")

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:

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"""

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/"

templates/about.html

@@ -41,31 +41,19 @@ <h2>What is BadgeTrack?</h2>
                         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.
-                    </p>                    <h2>How It Works</h2>
+                    </p>
+                    <h2>How It Works</h2>
                     <p>The application uses FastAPI and SQLite to:</p>
                     <ol>
-                        <li>Track visits to your specified URLs</li>
+                        <li>Track visits to your specified URLs using anonymous cookies</li>
                         <li>Generate dynamic badges via Shields.io integration</li>
                         <li>Provide real-time visit count updates</li>
                         <li>Offer customizable styles, colors, and logos</li>
-                        <li>Implement enhanced rate limiting for abuse prevention</li>
                     </ol>
 
-                    <h2>Rate Limiting</h2>
-                    <p>
-                        BadgeTrack implements a dual rate limiting system:
-                    </p>
-                    <ul>
-                        <li><strong>Visit Tracking:</strong> 48-hour cooldown per IP/URL combination to prevent spam</li>
-                        <li><strong>Badge Creation:</strong> Maximum 10 new badges per IP per day</li>
-                    </ul>
-
                     <h2>Privacy & Security</h2>
                     <p>
-                        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.
                     </p>
 
                     <h2>Open Source</h2>