From 71a4ad4a659d5159dd9aea40a7a61242becd1700 Mon Sep 17 00:00:00 2001 From: Elia El Lazkani Date: Sun, 13 Oct 2019 16:19:32 +0200 Subject: [PATCH] Migrated the ID generation system for security reasons. We have a new ID generator. This should make it hard to guess other user's IDs. This is configurable. --- .gitignore | 2 ++ config/config.yaml | 6 ++++ requirements/development.txt | 2 ++ shortenit/counter.py | 29 ---------------- shortenit/db.py | 25 +++++++------- shortenit/exceptions.py | 4 +++ shortenit/main.py | 64 +++++++++++++++++++++--------------- shortenit/pointer.py | 34 ++----------------- shortenit/shortener.py | 55 +++++++++++++++++++++++++++++++ shortenit/web.py | 6 ++-- 10 files changed, 127 insertions(+), 100 deletions(-) delete mode 100644 shortenit/counter.py create mode 100644 shortenit/exceptions.py create mode 100644 shortenit/shortener.py diff --git a/.gitignore b/.gitignore index 3422ba0..ef795fd 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ .eggs/ *.egg-info __pycache__/ +.mypy*/ +.vscode/ diff --git a/config/config.yaml b/config/config.yaml index e1ee101..4982854 100644 --- a/config/config.yaml +++ b/config/config.yaml @@ -5,3 +5,9 @@ CouchDB: username: root password: root url: http://localhost:5984 +Shortener: + # *CAUTION*: Enabling this check if the ID already exists before returning it. + # Even though this guarantees that the ID doesn't exist, this might inflict + # some performance hit. + check_duplicate_id: False + id_length: 32 \ No newline at end of file diff --git a/requirements/development.txt b/requirements/development.txt index 42e722f..e846edb 100644 --- a/requirements/development.txt +++ b/requirements/development.txt @@ -1,2 +1,4 @@ -r requirements.txt setuptools-git +setuptools-git-version +mypy diff --git a/shortenit/counter.py b/shortenit/counter.py deleted file mode 100644 index 4439415..0000000 --- a/shortenit/counter.py +++ /dev/null @@ -1,29 +0,0 @@ -import logging - -from cloudant.document import Document - - -class Counter: - def __init__(self, counter_db): - self.logger = logging.getLogger(self.__class__.__name__) - self.counter_db = counter_db - self.counter = None - - def get_counter(self) -> int: - with Document(self.counter_db, 'counter') as counter: - self.logger.debug("Counter: %s", counter) - try: - self.counter = counter['value'] - except KeyError: - self.logger.warn( - "Counter was not initialized, initializing...") - counter['value'] = 0 - try: - counter['value'] += 1 - except Exception as e: - self.logger.err(e) - # Need to check if the value exists or not as to not jump values - # which it currently does but it's not a big issue for right now - self.counter = counter['value'] - - return self.counter diff --git a/shortenit/db.py b/shortenit/db.py index 2390fa8..3d85795 100644 --- a/shortenit/db.py +++ b/shortenit/db.py @@ -1,7 +1,9 @@ import logging +import requests from cloudant.client import CouchDB +from shortenit.exceptions import DBConnectionFailed class DB: def __init__(self, config: dict) -> None: @@ -13,16 +15,6 @@ class DB: self.session = None def initialize_shortenit(self): - try: - self.counter_db = self.client['counter'] - except KeyError: - self.logger.warn( - "The 'counter' database was not found, creating...") - self.counter_db = self.client.create_database('counter') - if self.counter_db.exists(): - self.logger.info( - "The 'counter' database was successfully created.") - try: self.data_db = self.client['data'] except KeyError: @@ -46,8 +38,17 @@ class DB: def __enter__(self) -> CouchDB: """ """ - self.client = CouchDB(self.username, self.password, - url=self.url, connect=True) + try: + self.client = CouchDB(self.username, self.password, + url=self.url, connect=True) + except requests.exceptions.ConnectionError as e: + self.logger.fatal("Failed to connect to database, is it on?") + self.logger.fatal("%s", e) + raise DBConnectionFailed + except requests.exceptions.HTTPError as e: + self.logger.fatal("Failed to authenticate to database.") + self.logger.fatal("%s", e) + raise DBConnectionFailed self.session = self.client.session() return self diff --git a/shortenit/exceptions.py b/shortenit/exceptions.py new file mode 100644 index 0000000..2033f1d --- /dev/null +++ b/shortenit/exceptions.py @@ -0,0 +1,4 @@ +import sys + +class DBConnectionFailed(Exception): + pass \ No newline at end of file diff --git a/shortenit/main.py b/shortenit/main.py index 00ed060..4289a6f 100644 --- a/shortenit/main.py +++ b/shortenit/main.py @@ -1,4 +1,5 @@ #!/usr/bin/env python3 +import sys import argparse import logging import pathlib @@ -6,13 +7,14 @@ import asyncio import time -from .data import Data -from .pointer import Pointer -from .config import Config -from .counter import Counter -from .db import DB -from .logger import setup_logging -from .web import Web, SiteHandler +from shortenit.data import Data +from shortenit.pointer import Pointer +from shortenit.config import Config +from shortenit.shortener import Shortener +from shortenit.db import DB +from shortenit.logger import setup_logging +from shortenit.web import Web, SiteHandler +from shortenit.exceptions import DBConnectionFailed PROJECT_ROOT = pathlib.Path(__file__).parent.parent CONFIGURATION = f'{PROJECT_ROOT}/config/config.yaml' @@ -34,28 +36,34 @@ def main() -> None: db_config = config.get('CouchDB', None) server_config = config.get('Server', None) if db_config: - with DB(db_config) as db: - db.initialize_shortenit() + try: + with DB(db_config) as db: + db.initialize_shortenit() - handler = SiteHandler(db, shorten_url, lenghten_url) - web = Web(handler, debug=debug) - web.host = server_config.get('host', None) - web.port = server_config.get('port', None) - web.start_up() + handler = SiteHandler(config, db, shorten_url, lenghten_url) + web = Web(handler, debug=debug) + web.host = server_config.get('host', None) + web.port = server_config.get('port', None) + web.start_up() + except DBConnectionFailed as e: + sys.exit(1) + sys.exit(0) -def shorten_url(database: DB, data: str, ttl: time.time): - counter = Counter(database.counter_db) - data = Data(database.data_db, - data=data) - data.populate() - pointer = Pointer(database.pointers_db, counter) - pointer.generate_pointer( - data.identifier, - ttl - ) - data.set_data(pointer.identifier) - return pointer.identifier +def shorten_url(configuration: dict, database: DB, + data: str, ttl): + shortener = Shortener(database.pointers_db, + configuration.get('Shortener', None)) + identifier = shortener.get_id() + if identifier: + _data = Data(database.data_db, + data=data) + _data.populate() + pointer = Pointer(database.pointers_db, identifier) + pointer.generate_pointer(_data.identifier, ttl) + _data.set_data(pointer.identifier) + return pointer.identifier + return None def lenghten_url(database: DB, identifier: str): @@ -100,3 +108,7 @@ def verbosity(verbose: int): return logging.INFO elif verbose > 2: return logging.DEBUG + + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/shortenit/pointer.py b/shortenit/pointer.py index 65bd09b..76dbd7c 100644 --- a/shortenit/pointer.py +++ b/shortenit/pointer.py @@ -1,7 +1,6 @@ import time import logging -from .counter import Counter from cloudant.document import Document CHARS = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ' @@ -9,21 +8,16 @@ CHARS = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ' class Pointer: def __init__(self, pointers_db: object, - counter: Counter = None) -> None: + identifier: str = None) -> None: self.logger = logging.getLogger(self.__class__.__name__) self.pointers_db = pointers_db - self.counter = counter - self.identifier = None + self.identifier = identifier self.data_hash = None self.ttl = None self.timestamp = time.time() def generate_pointer(self, data_hash: str, ttl: time.time): - self.logger.debug("Generating new counter...") - counter = self.counter.get_counter() - self.logger.debug("Encoding the counter into an ID") - self.identifier = Pointer.encode(counter) - self.logger.debug("Encoded counter is %s", self.identifier) + self.logger.debug("identifier is %s", self.identifier) with Document(self.pointers_db, self.identifier) as pointer: pointer['value'] = data_hash pointer['ttl'] = ttl @@ -43,25 +37,3 @@ class Pointer: except KeyError: pass return None - - @staticmethod - def encode(counter): - sign = '-' if counter < 0 else '' - counter = abs(counter) - result = '' - while counter > 0: - counter, remainder = divmod(counter, len(CHARS)) - result = CHARS[remainder]+result - - return sign+result - - @staticmethod - def decode(counter): - return int(counter, len(CHARS)) - - @staticmethod - def padding(counter: str, count: int=6): - if len(counter) < count: - pad = '0' * (count - len(counter)) - return f"{pad}{counter}" - return f"{counter}" diff --git a/shortenit/shortener.py b/shortenit/shortener.py new file mode 100644 index 0000000..8081421 --- /dev/null +++ b/shortenit/shortener.py @@ -0,0 +1,55 @@ +import uuid +import logging + +from cloudant.document import Document + + +class Shortener: + def __init__(self, pointer_db, configuration: dict): + self.logger = logging.getLogger(self.__class__.__name__) + self.pointer_db = pointer_db + self.uuid = None + self.length = 32 + self.check_duplicate = False + self.configuration = configuration + self.init() + + def init(self): + length = self.configuration.get('id_length', 32) + if length >= 32 or length <= 0: + self.length = 32 + else: + self.length = length + self.check_duplicate = self.configuration.get( + 'check_duplicate_id', False) + + def generate_short_uuid(self): + short_uuid = uuid.uuid1().hex + return short_uuid.upper()[0:self.length] + + def check_uuid(self, short_uuid): + with Document(self.pointer_db, 'pointer') as pointer: + self.logger.debug("Pointer: %s", pointer) + try: + self.uuid = pointer[short_uuid] + except KeyError: + self.logger.info("Generated short uuid '%s'" + "was not found in database", + short_uuid) + return False + return True + + def get_id(self): + short_uuid = self.generate_short_uuid() + if self.check_duplicate: + counter = 0 + while self.check_uuid(short_uuid): + if counter > 10: + self.logger.err("Cannot generate new unique ID," + "try to configure a longer ID length.") + return None + short_uuid = self.generate_short_uuid() + counter += 1 + self.logger.debug("Returning ID: '%s'", short_uuid) + return short_uuid + diff --git a/shortenit/web.py b/shortenit/web.py index 2b8c226..ba4c8ba 100644 --- a/shortenit/web.py +++ b/shortenit/web.py @@ -38,8 +38,9 @@ class Web: class SiteHandler: - def __init__(self, database, shorten_url, lenghten_url): + def __init__(self, configuration, database, shorten_url, lenghten_url): self.logger = logging.getLogger(self.__class__.__name__) + self.configuration = configuration self.database = database self.shorten_url = shorten_url self.lenghten_url = lenghten_url @@ -59,7 +60,8 @@ class SiteHandler: abort(400) try: short_url = self.shorten_url( - self.database, data['url'], data['timestamp']) + self.configuration, self.database, + data['url'], data['timestamp']) except KeyError as e: self.logger.error(e) abort(400)