From 1836ccbc6b416b35c8b266f4050e76eedb80df64 Mon Sep 17 00:00:00 2001 From: Robert Putt Date: Thu, 30 Nov 2017 20:16:19 +0000 Subject: [PATCH] Add basic database model Adds basic database model to get us started. This will need to be extended drastically for basic functionality but should allow us to make a simple worker that simply adds hashes to the DB. Long term maybe MySQL isnt even a good idea, but it'll get a basic proof of concept sorted for now. Change-Id: Ifdcd2309fa2ad6a9a27dd6c3ba61dd8db0dcb854 --- alembic/env.py | 46 +++++++++++++++++++---------------- nemesis.conf.sample | 1 - python_nemesis/db/__init__.py | 0 python_nemesis/db/models.py | 41 +++++++++++++++++++++++++++++++ 4 files changed, 66 insertions(+), 22 deletions(-) create mode 100644 python_nemesis/db/__init__.py create mode 100644 python_nemesis/db/models.py diff --git a/alembic/env.py b/alembic/env.py index c1b0b20..c43f079 100644 --- a/alembic/env.py +++ b/alembic/env.py @@ -12,29 +12,31 @@ # License for the specific language governing permissions and limitations # under the License. +''' + Alembic environment script as generated by alembic init, modified to + use the Flask app context and Flask-SQLAlchemy extension rather than + the traditional raw SQLAlchemy engine. +''' + + from __future__ import with_statement from alembic import context -from sqlalchemy import engine_from_config, pool +import configparser from logging.config import fileConfig +import os +from python_nemesis.extensions import db +from sqlalchemy import engine_from_config +from sqlalchemy import pool + -# this is the Alembic Config object, which provides -# access to the values within the .ini file in use. config = context.config - -# Interpret the config file for Python logging. -# This line sets up loggers basically. fileConfig(config.config_file_name) +target_metadata = db.metadata +config_file = os.environ.get("NEMESIS_CONFIG", + "/etc/nemesis/nemesis.conf") -# add your model's MetaData object here -# for 'autogenerate' support -# from myapp import mymodel -# target_metadata = mymodel.Base.metadata -target_metadata = None - -# other values from the config, defined by the needs of env.py, -# can be acquired: -# my_important_option = config.get_main_option("my_important_option") -# ... etc. +nem_config = configparser.ConfigParser() +nem_config.readfp(open(config_file)) def run_migrations_offline(): @@ -49,7 +51,7 @@ def run_migrations_offline(): script output. """ - url = config.get_main_option("sqlalchemy.url") + url = nem_config.get('sqlalchemy', 'database_uri') context.configure( url=url, target_metadata=target_metadata, literal_binds=True) @@ -64,10 +66,11 @@ def run_migrations_online(): and associate a connection with the context. """ - connectable = engine_from_config( - config.get_section(config.config_ini_section), - prefix='sqlalchemy.', - poolclass=pool.NullPool) + db_uri = nem_config.get('sqlalchemy', 'database_uri') + config_dict = {"sqlalchemy.url": db_uri} + connectable = engine_from_config(config_dict, + prefix='sqlalchemy.', + poolclass=pool.NullPool) with connectable.connect() as connection: context.configure( @@ -78,6 +81,7 @@ def run_migrations_online(): with context.begin_transaction(): context.run_migrations() + if context.is_offline_mode(): run_migrations_offline() else: diff --git a/nemesis.conf.sample b/nemesis.conf.sample index 362b21e..a946da5 100644 --- a/nemesis.conf.sample +++ b/nemesis.conf.sample @@ -13,4 +13,3 @@ delay_auth_decision = True [sqlalchemy] database_uri = sqlite:////tmp/nemesis.db echo = true - diff --git a/python_nemesis/db/__init__.py b/python_nemesis/db/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python_nemesis/db/models.py b/python_nemesis/db/models.py new file mode 100644 index 0000000..f796f57 --- /dev/null +++ b/python_nemesis/db/models.py @@ -0,0 +1,41 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); you may +# not use this file except in compliance with the License. You may obtain +# a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations +# under the License. + +from python_nemesis.extensions import db + + +class Files(db.Model): + __tablename__ = 'file' + file_id = db.Column(db.BigInteger(), primary_key=True, autoincrement=True) + sha512_hash = db.Column(db.UnicodeText(), nullable=True, index=True) + sha256_hash = db.Column(db.UnicodeText(), nullable=True, index=True) + sha1_hash = db.Column(db.UnicodeText(), nullable=True, index=True) + md5_hash = db.Column(db.UnicodeText(), nullable=True, index=True) + crc32 = db.Column(db.UnicodeText(), nullable=True, index=True) + size = db.Column(db.Float(), nullable=True) + mime_type = db.Column(db.String(40), nullable=True) + submitted_by = db.Column(db.String(120), nullable=False, index=True) + status = db.Column(db.String(20), nullable=False) + last_updated = db.Column(db.DateTime, nullable=False) + first_seen = db.Column(db.DateTime, nullable=False) + role_function = db.relationship("FileLookupRequest") + + +class FileLookupRequest(db.Model): + __tablename__ = 'lookup_request' + request_id = db.Column(db.BigInteger(), primary_key=True, + autoincrement=True) + requested_at = db.Column(db.DateTime, nullable=False) + requestor = db.Column(db.String(120), nullable=False) + file_id = db.Column(db.ForeignKey(Files.file_id), nullable=True) + lookup_hash = db.Column(db.UnicodeText(), nullable=False) + result = db.Column(db.String(20), nullable=False)