summaryrefslogtreecommitdiff
path: root/fbin/db.py
blob: 2bf153bf68522fd0812f22cf696c60abe74225d8 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
from contextlib import contextmanager
import datetime
import mimetypes
import os

from flask import current_app
from sqlalchemy import create_engine, Column, Integer, String, DateTime, Text, Index, ForeignKey, Boolean, JSON
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, relation, backref
from sqlalchemy.orm.exc import NoResultFound
from sqlalchemy.exc import IntegrityError
from sqlalchemy.sql import and_

engine = create_engine(current_app.config['DB_URI'])

Base = declarative_base(bind = engine)

class User(Base):
    __tablename__ = 'users'

    id = Column(Integer, primary_key = True)
    username = Column(String, unique = True, index = True)
    jab_id = Column(String(24), unique = True, index = True)
    api_key_date = Column(DateTime, default = datetime.datetime.utcnow)
    files = relation('File', backref = 'user', order_by = 'File.date.desc()')

    def __init__(self, username, jab_id):
        self.username = username
        self.jab_id = jab_id

class UserSession(Base):
    __tablename__ = 'sessions'

    id = Column(Integer, primary_key = True)
    user_id = Column(Integer, ForeignKey('users.id'), index = True)
    access_token = Column(String)
    refresh_token = Column(String)
    updated = Column(DateTime)

    def __init__(self, user_id, access_token, refresh_token):
        self.user_id = user_id
        self.access_token = access_token
        self.refresh_token = refresh_token

class File(Base):
    __tablename__ = 'files'

    id = Column(Integer, primary_key = True)
    hash = Column(String, unique = True, index = True)
    filename = Column(String)
    date = Column(DateTime)
    user_id = Column(Integer, ForeignKey('users.id'), nullable = True)
    ip = Column(String)
    accessed = Column(DateTime)
    scanned = Column(Boolean, nullable=False, default=False)
    blocked_reason = Column(JSON)

    def __init__(self, hash, filename, date, user_id = None, ip = None):
        self.hash = hash
        self.filename = filename
        self.date = date
        self.user_id = user_id
        self.ip = ip

    @staticmethod
    def pretty_size(size):
        if size is None:
            return 'N/A'
        suffixes = (('TiB', 2**40), ('GiB', 2**30), ('MiB', 2**20), ('KiB', 2**10))
        for suf, threshold in suffixes:
            if size >= threshold:
                return '{:.2f} {}'.format(size / threshold, suf)
            else:
                continue
        return '{} B'.format(size)

    def get_path(self):
        return os.path.join(current_app.config['FILE_DIRECTORY'], self.hash + os.path.splitext(self.filename)[1])

    def get_thumb_path(self):
        return os.path.join(current_app.config['THUMB_DIRECTORY'], self.hash + '.jpg')

    def get_size(self):
        try:
            return os.path.getsize(self.get_path())
        except OSError:
            return None

    @property
    def formatted_size(self):
        return self.pretty_size(self.get_size())

    @property
    def formatted_date(self):
        return self.date.strftime('%Y-%m-%d %H:%M:%S UTC')

    def get_mime_type(self):
        return mimetypes.guess_type(self.filename, strict = False)[0] or 'application/octet-stream'

    def is_image(self):
        return self.get_mime_type().startswith('image')

    def is_video(self):
        return self.get_mime_type().startswith('video')

    @property
    def ext(self):
        return os.path.splitext(self.filename)[1]

    @property
    def exists(self):
        return os.path.exists(self.get_path())

Base.metadata.create_all()
Session = sessionmaker(bind = engine, autoflush = True, autocommit = False)

@contextmanager
def session_scope():
    session = Session()
    try:
        session.expire_on_commit = False
        yield session
        session.commit()
    except:
        session.rollback()
        raise
    finally:
        session.close()