def get_or_create(Model, **kwargs):
'''
A get_or_create method exactly like the peewee's one, but compatible with
SQLite (does not starts a transaction if SQLite is used because it does not
support nested transactions).
'''
defaults = kwargs.pop('defaults', {})
query = Model.select()
for field, value in kwargs.items():
if '__' in field:
query = query.filter(**{field: value})
else:
query = query.where(getattr(Model, field) == value)
try:
return query.get(), False
except Model.DoesNotExist:
try:
params = dict((k, v) for k, v in kwargs.items() if '__' not in k)
params.update(defaults)
if type(Model._meta.database) == SqliteDatabase:
return Model.create(**params), True
else:
with Model._meta.database.atomic():
return Model.create(**params), True
except IntegrityError as exc:
try:
return query.get(), False
except Model.DoesNotExist:
raise exc
def _descriptor_set(self: peewee.FieldDescriptor, instance: Model, value: Any) -> None:
"""Monkeypatch the __set__ method on peewee descriptors to always save immediately.
This is useful for us because in interactive use, it is easy to forget to call .save(), and
we are not concerned about the performance implications of saving often.
"""
instance._data[self.att_name] = value
instance._dirty.add(self.att_name)
# Otherwise this gets called in the constructor.
if getattr(instance, "_is_prepared", False):
instance.save()
def _get_model_value(model: Model,
cascade: Union[bool, int],
null: bool = False,
**filters) -> Union[dict, int]:
"""Converts a model to a JSON value."""
cascade = _check_cascade(cascade)
if cascade.cascade:
with suppress(AttributeError):
return model.to_json(null=null, cascade=cascade.next, **filters)
return model.get_id()
def _load_table(table: Model, directory: Path, format_: str):
"""Dump a single table."""
if directory is not None:
print(f" Loading {table.table_name()}...")
in_file = Path(directory) / f"{table.table_name()}.{format_}"
dataset = tablib.Dataset(headers=table.fields()).load(
in_file.read_text())
print(f" Importing {table.table_name()} into the database...")
table.insert_many(dataset.dict).execute()
print(" Done.")
print("=====================")
else:
pass
def save(self, return_instance=True, **kwargs):
with database_proxy.transaction():
rows = BaseModel.save(self, **kwargs)
if return_instance:
pk = self._get_pk_value()
return self.get(id=pk)
return rows
def save(self, return_instance=True, **kwargs):
with database_proxy.transaction():
rows = BaseModel.save(self, **kwargs)
if return_instance:
pk = self._get_pk_value()
return self.get(id=pk)
return rows
def createTable(Model, db_name, user=None, password=None, host='127.0.0.1'):
'''
create by bigzhu at 15/04/04 01:08:30 建立数据库表; peewee 要在定义时候指定 db 相当不人性化,修正
modify by bigzhu at 15/04/04 01:32:46 没有这个数据库的时候,直接返回建立数据的语句
modify by bigzhu at 15/04/04 01:45:43 如果表已经存在,不能往下继续了
'''
if user is None:
user = db_name
if password is None:
password = db_name
if host is None:
host = '127.0.0.1'
#db = PostgresqlExtDatabase(db_name, user=user, password=password, host='127.0.0.1', register_hstore=False)
db = PostgresqlExtDatabase(db_name, user=user, password=password, host=host, register_hstore=False)
Model._meta.database = db
try:
if Model.table_exists():
print 'table %s already exists' % Model.__name__
return
createBaseTable(db)
Model.create_table()
print 'create table ' + Model.__name__
except peewee.OperationalError:
print public_bz.getExpInfo()
showDBCreate(db_name)
exit(1)
table_name = Model.__name__
if table_name != 'base':
sql = '''
alter table %s inherit base;
''' % table_name
db.execute_sql(sql)
resetBaseDefault(db)
# add table comment
comment = Model.__doc__
sql = '''
COMMENT ON TABLE %s IS '%s';
''' % (table_name, comment)
db.execute_sql(sql)
def dropTable(Model, db_name, user=None, password=None, host='127.0.0.1'):
'''
create by bigzhu at 15/04/04 13:12:02 还是需要一个删除表的功能
'''
if user is None:
user = db_name
if password is None:
password = db_name
if host is None:
host = '127.0.0.1'
#db = PostgresqlDatabase(db_name, user=user, password=password, host='127.0.0.1')
#db = PostgresqlExtDatabase(db_name, user=user, password=password, host='127.0.0.1', register_hstore=False)
db = PostgresqlExtDatabase(db_name, user=user, password=password, host=host, register_hstore=False)
Model._meta.database = db
try:
Model.drop_table(True)
except peewee.OperationalError:
print public_bz.getExpInfo()
showDBCreate(db_name)
exit(1)
print 'drop table ' + Model.__name__
def save(self, **kwargs):
with vulner_db_proxy.transaction():
self.cvss_rank = floor(float(self.cvss_score))
if isinstance(self.published, datetime):
self.published = dt2str(self.published)
if isinstance(self.modified, datetime):
self.modified = dt2str(self, modified)
if isinstance(self.last_seen, datetime):
self.last_seen = dt2str(self.last_seen)
if isinstance(self.created, datetime):
self.created = dt2str(self.created)
if isinstance(self.cvss_time, datetime):
self.cvss_time = dt2str(self.cvss_time)
self.cwe_id_list = [onlydigits(cwe) for cwe in self.cwe_list]
self.componentversions_string = ";".join(self.componentversions)
# TODO: Check if self.id defined after Model.save or before Model.save
if self.vulner_id == default_vulner_id_undefined:
self.vulner_id = default_vulner_id_start + self.component + default_vulner_id_delimeter + self.id
Model.save(self, **kwargs)
def get_pages(
websites: tp.Iterable[tp.Iterable[str]],
session: _Session,
token_model: peewee.Model,
timeout: int = 5) -> tp.Iterable[tp.Tuple[str, tp.List[str]]]:
"""Fetches page urls from websites' main pages and stores their's tokens."""
Controller.set_page_loading_state(0, len(websites), 0)
for i, website in enumerate(websites):
website_url, page_pattern, *adt = website
try:
if not (res := session.get(website_url,
timeout=(timeout, timeout))).ok:
continue
for j, url in enumerate(
re.findall(page_pattern, res.content.decode("utf-8"))):
url = Loader._expand_url(url, website_url)
if not token_model.get_or_none(token_model.token == (
token := Loader._encode_url(url))):
token_model.create(token=token).save()
yield url, adt
Controller.set_page_loading_state(
i + 1, len(websites), j + 1)
def _dump_table(table: Model, directory: Path, format_: str):
"""Dump a single table."""
try:
table.select().tuples()
table.fields()
dataset = tablib.Dataset(*table.select().tuples(),
headers=table.fields())
except:
print(table._meta.database.get_columns(table.table_name()))
if directory is not None:
print(f" Dumping {table.table_name()}...")
out_file = Path(directory) / f"{table.table_name()}.{format_}"
out_file.write_text(dataset.export(format_))
print(" Done.")
print("=====================")
else:
print(dataset.export("csv"))
def delete_user_from_table(self, Model, user):
'''
Deletes a given user's records from a table in preparation for overwriting.
'''
Model.delete().where(Model.user == user.uuid).execute()
def get_without_failing(Model, query):
results = Model.select().where(query).limit(1)
return results[0] if len(results) > 0 else None
def save(self, **kwds):
with db.transaction() as txn:
Model.save(self, **kwds)
def _add_api_keys_to_users_table(table: Model, _column: Field) -> None:
log.info('Adding API Keys for all users, might take some extra time...')
with db_config.database.transaction():
for user in table:
user.api_key = table.random_password(stronger=True)
user.save()
def __init__(self, **kwargs):
if self.table_exists() is False:
self.create_table()
Model.__init__(self, **kwargs)
def save(self, obj: Model, save: bool = True) -> Model:
self.populate_obj(obj)
if save:
obj.save()
return obj
