From 71f4785d13aa51b2858f182cf2dd25bb93996661 Mon Sep 17 00:00:00 2001
From: Jeremy Zhang <jeremyzhang5000@gmail.com>
Date: Fri, 21 Apr 2017 18:06:19 -0700
Subject: [PATCH] Initial Redislite implementation with limits

---
 .gitignore                               |  2 ++
 titanembeds/app.py                       |  1 +
 titanembeds/database/__init__.py         |  1 +
 titanembeds/database/custom_redislite.py | 25 ++++++++++++++++++++++++
 4 files changed, 29 insertions(+)
 create mode 100644 titanembeds/database/custom_redislite.py

diff --git a/.gitignore b/.gitignore
index 09526cb..591a53a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -90,3 +90,5 @@ ENV/
 
 # Project specifc
 config.py
+redislite.db
+redislite.db.settings
diff --git a/titanembeds/app.py b/titanembeds/app.py
index 6d5b2d2..90a506c 100644
--- a/titanembeds/app.py
+++ b/titanembeds/app.py
@@ -15,6 +15,7 @@ app.config['SQLALCHEMY_DATABASE_URI'] = config['database-uri']
 app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False  # Suppress the warning/no need this on for now.
 app.config['RATELIMIT_HEADERS_ENABLED'] = True
 app.config['SQLALCHEMY_POOL_RECYCLE'] = 250
+app.config['RATELIMIT_STORAGE_URL'] = 'redislite://redislite.db'
 app.secret_key = config['app-secret']
 
 db.init_app(app)
diff --git a/titanembeds/database/__init__.py b/titanembeds/database/__init__.py
index 1fb2b67..28951e6 100644
--- a/titanembeds/database/__init__.py
+++ b/titanembeds/database/__init__.py
@@ -5,3 +5,4 @@ from guilds import Guilds
 from unauthenticated_users import UnauthenticatedUsers
 from unauthenticated_bans import UnauthenticatedBans
 from authenticated_users import AuthenticatedUsers
+from custom_redislite import LimitsRedisLite
diff --git a/titanembeds/database/custom_redislite.py b/titanembeds/database/custom_redislite.py
new file mode 100644
index 0000000..e79849c
--- /dev/null
+++ b/titanembeds/database/custom_redislite.py
@@ -0,0 +1,25 @@
+import urlparse
+from limits.storage import Storage
+from redislite import Redis
+
+class LimitsRedisLite(Storage): # For Python Limits
+    STORAGE_SCHEME = "redislite"
+    def __init__(self, uri, **options):
+        self.redis_instance = Redis(urlparse.urlparse(uri).netloc)
+
+    def check(self):
+        return True
+
+    def get_expiry(self, key):
+        return self.redis_instance.ttl(key)
+
+    def incr(self, key, expiry, elastic_expiry=False):
+        if self.redis_instance.exists(key):
+            self.redis_instance.set(key, int(self.redis_instance.get(key))+1)
+        else:
+            self.redis_instance.set(key, 1)
+        self.redis_instance.expire(key, expiry)
+        return
+
+    def get(self, key):
+        return int(self.redis_instance.get(key))