class SubJobResult(db.Model):
STATUS_OK = 1
STATUS_DELETE = 2
id = db.Column(db.BigInteger, primary_key=True)
sub_job_id = db.Column(db.BigInteger, db.ForeignKey('sub_job.id'))
result = db.Column(db.Text)
created_time = db.Column(db.DateTime)
status = db.Column(db.Integer, nullable=False)
def __init__(self, sub_job, result, created_time=None, status=None):
self.sub_job_id = sub_job.id
self.result = json.dumps(result) if isinstance(result, (
dict,
list,
tuple,
)) else result
self.created_time = timezone.now(
) if created_time is None else created_time
self.status = self.STATUS_OK if status is None else status
@classmethod
def create(cls, sub_job, result):
job_result = cls(sub_job=sub_job, result=result)
db.session.add(job_result)
db.session.commit()
return job_result
class Role(db.Model):
id = db.Column(db.Integer, primary_key=True)
name = db.Column(db.String(32), nullable=True, default='')
type = db.Column(db.Integer, nullable=True, default=0)
created_time = db.Column(db.DateTime, nullable=False)
status = db.Column(db.Integer, nullable=False, default=0)
users = db.relationship('User', backref='role', lazy='dynamic')
class AssetSysVul(db.Model, AsDictMixin):
id = db.Column(db.BigInteger, primary_key=True)
ip = db.Column(db.String(256), nullable=False)
plugin_ident = db.Column(db.String(128), nullable=False)
payloads = db.Column(db.Text)
created_time = db.Column(db.DateTime, nullable=False)
last_discover_time = db.Column(db.DateTime, nullable=False)
@classmethod
def stats(cls):
return {
'total':
cls.query.count(),
'24_hour':
cls.query.filter(cls.created_time >= timezone.now() -
timedelta(days=1)).count(),
}
@classmethod
def stats_host_vul(cls, topn=10):
rs = db.session.execute(
'select ip, count(*) from asset_sys_vul group by ip order by count desc limit :topn',
{'topn': topn})
return [dict(r) for r in rs]
@classmethod
def stats_vul_host(cls, topn=10):
rs = db.session.execute(
'select plugin_ident, count(*) from asset_sys_vul group by plugin_ident order by count desc limit :topn',
{'topn': topn})
return [dict(r) for r in rs]
@classmethod
def create_or_replace(cls, key, plugin_ident, payloads):
obj = cls.query.filter_by(ip=key, plugin_ident=plugin_ident).first()
if obj is None:
obj = cls()
obj.ip = key
obj.plugin_ident = plugin_ident
obj.created_time = timezone.now()
obj.last_discover_time = timezone.now()
obj.payloads = json.dumps(payloads) if isinstance(
payloads, (list, dict)) else payloads
db.session.add(obj)
db.session.commit()
@classmethod
def delete_not_found(cls, key, now):
cls.query.filter(cls.ip == key, cls.last_discover_time < now).delete()
@classmethod
def all(cls, key=None):
query = cls.query
if key:
query = query.filter_by(ip=key)
return query.all()
class WebAsset(db.Model):
STATUS_NEW = 0
STATUS_ACK = 1
STATUS_OFFLINE = 2
STATUS_DELETE = 3
id = db.Column(db.BigInteger, primary_key=True)
name = db.Column(db.String(32), nullable=False, default='')
url = db.Column(db.String(256), nullable=False, default='')
group_id = db.Column(db.BigInteger, db.ForeignKey('web_group.id'))
sys_asset_id = db.Column(db.BigInteger, db.ForeignKey('sys_asset.id'))
manager = db.Column(db.String(64), nullable=False, default='')
created_time = db.Column(db.DateTime, nullable=False, default='')
last_discove_time = db.Column(db.DateTime, nullable=False, default='')
status = db.Column(db.Integer, nullable=False)
@classmethod
def all(cls):
return cls.query.filter(cls.status.notin_([cls.STATUS_DELETE])).all()
class WebGroup(db.Model):
STATUS_OK = 0
STATUS_DELETE = 1
id = db.Column(db.BigInteger, primary_key=True)
pid = db.Column(db.BigInteger, nullable=False)
name = db.Column(db.String(32), nullable=False, default='')
created_time = db.Column(db.DateTime, nullable=False)
manager = db.Column(db.String(64), nullable=False, default='')
status = db.Column(db.Integer, nullable=False)
assets = db.relationship('WebAsset', backref="group", lazy="dynamic")
class SysAsset(db.Model, AsDictMixin):
WEIGHT_HIGH = 10
WEIGHT_MEDIUM = 6
WEIGHT_LOW = 1
STATUS_NEW = 0
STATUS_ACK = 1
STATUS_DELETE = 2
id = db.Column(db.BigInteger, primary_key=True)
sn = db.Column(db.String(128), nullable=False, default='')
name = db.Column(db.String(32), nullable=False, default='')
ip = db.Column(db.String(256), nullable=False, default='')
mac = db.Column(db.String(32), nullable=False, default='')
os = db.Column(db.String(128), nullable=False, default='')
group_id = db.Column(db.BigInteger, db.ForeignKey('sys_group.id'))
manager = db.Column(db.String(64), nullable=False, default='')
weight = db.Column(db.Integer, nullable=False)
created_time = db.Column(db.DateTime, nullable=False)
last_discove_time = db.Column(db.DateTime, nullable=False)
status = db.Column(db.Integer, nullable=False)
applications = db.relationship('Application', backref='asset', lazy='dynamic')
@classmethod
def stats(cls):
return {
'total' : cls.query.filter(cls.status.notin_([cls.STATUS_DELETE])).count(),
'24_hour' : cls.query.filter(cls.created_time>=timezone.now() - timedelta(days=1), cls.status.notin_([cls.STATUS_DELETE])).count(),
}
@classmethod
def all(cls):
return cls.query.filter(cls.status.notin_([cls.STATUS_DELETE])).all()
@classmethod
def get_by_key(cls, value, key='id'):
return db.session.query(cls).filter_by(**{key : value}).first()
@classmethod
def create_or_replace(cls, **kwargs):
ip = kwargs.get('ip', '')
obj = cls.get_by_key(ip, key='ip')
try:
if obj is None or obj.status == cls.STATUS_DELETE:
obj = cls()
obj.ip = ip
obj.status = cls.STATUS_NEW
obj.created_time = timezone.now()
obj.weight = cls.WEIGHT_LOW
obj.last_discove_time = timezone.now()
obj.name = kwargs.get('name', obj.name)
obj.os = kwargs.get('os', obj.os)
obj.mac = kwargs.get('mac', obj.mac)
db.session.add(obj)
db.session.commit()
for app in kwargs.get('apps', []):
Application.create_or_replace(app.get('name', ''),
app.get('protocol', ''),
app.get('port', 0),
app.get('state', ''),
app.get('product', ''),
app.get('version', ''),
obj,
)
except BaseException as e:
logger.error(e)
logger.exception(traceback.format_exc())
db.session.rollback()
return obj
@classmethod
def delete_by_key(cls, id):
obj = cls.get_by_key.get(id)
if obj:
obj.status = cls.STATUS_DELETE
db.session.add(obj)
db.session.commit()
return obj
class Application(db.Model):
STATUS_NEW = 0
STATUS_ACK = 1
STATUS_DELETE = 2
id = db.Column(db.BigInteger, primary_key=True)
name = db.Column(db.String(128), nullable=False, default='')
protocol = db.Column(db.String(32), nullable=False, default='')
port = db.Column(db.Integer, nullable=False, default=0)
state = db.Column(db.String(32), nullable=False, default='')
product = db.Column(db.String(128), nullable=False, default='')
version = db.Column(db.String(64), nullable=False, default='')
asset_id = db.Column(db.BigInteger, db.ForeignKey('sys_asset.id'))
created_time = db.Column(db.DateTime, nullable=False, default='')
last_discove_time = db.Column(db.DateTime, nullable=False, default='')
status = db.Column(db.Integer, nullable=False)
@classmethod
def stats(cls):
return {
'total' : cls.query.filter(cls.status.notin_([cls.STATUS_DELETE])).count(),
'24_hour' : cls.query.filter(cls.created_time>=timezone.now() - timedelta(days=1), cls.status.notin_([cls.STATUS_DELETE])).count(),
}
@classmethod
def stats_port(cls):
rs = db.session.execute('select port, count(*) from application where status != :status group by port order by count desc', {'status' : cls.STATUS_DELETE})
return [dict(r) for r in rs]
@classmethod
def all(cls):
return cls.query.filter(cls.status.notin_([cls.STATUS_DELETE])).all()
@classmethod
def create_or_replace(cls, name, protocol, port, state, product, version, asset, *args, **kwargs):
obj = db.session.query(cls).filter_by(asset_id=asset.id, port=port).first()
if obj is None or obj.status == cls.STATUS_DELETE:
obj = cls()
obj.port = port
obj.created_time = timezone.now()
obj.status = cls.STATUS_NEW
obj.asset_id = asset.id
obj.name = name
obj.protocol = protocol
obj.state = state
obj.product = product
obj.last_discove_time = timezone.now()
db.session.add(obj)
db.session.commit()
class SysVulPlugin(db.Model, ValidatorMixin, AsDictMixin):
STATUS_OK = 1
STATUS_DELETE = 2
LEVEL_HEIGHT = 10
LEVEL_MEDIUM = 6
LEVEL_LOW = 1
TYPE_SCRIPT = 'script'
TYPE_SOCKET = 'socket_request'
TYPE_HTTP = 'http_request'
id = db.Column(db.BigInteger, primary_key=True)
ident = db.Column(db.String(128), nullable=False, default='')
name = db.Column(db.String(128), nullable=False, default='')
level = db.Column(db.Integer, nullable=False)
type = db.Column(db.String(128), nullable=False, default='')
params = db.Column(db.Text)
remark = db.Column(db.Text)
links = db.Column(db.Text)
created_time = db.Column(db.DateTime, nullable=False)
status = db.Column(db.Integer, nullable=False)
def __init__(self,
ident=None,
name=None,
level=None,
type=None,
params=None,
remark=None,
links=None,
created_time=None,
status=None,
id=None,
from_cache=False):
self.id = id
self.ident = ident
self.name = name
self.level = level
self.type = type
self.params = params
self.remark = remark
self.links = links
self.created_time = timezone.now(
) if created_time is None else created_time
self.status = self.STATUS_OK if status is None else status
@property
def params_object(self):
try:
return json.loads(self.params)
except BaseException as e:
logger.exception(e)
logger.error(traceback.format_exc())
return {}
@classmethod
def all(cls):
return db.session.query(cls).filter(
cls.status != cls.STATUS_DELETE).all()
@classmethod
def get_by_key(cls, value, key='id', all=False):
params = {key: value}
if not all:
params['status'] = cls.STATUS_OK
return cls.query.filter_by(**params).first()
def to_cache(self, clear=False):
if clear:
redis.delete(REDIS_KEYS['PLUGIN_SYS_VUL'].format(ident=self.ident))
else:
redis.hmset(REDIS_KEYS['PLUGIN_SYS_VUL'].format(ident=self.ident),
self.as_dict_string())
return True
@classmethod
def get_by_cache(cls, ident):
bplugin = redis.hgetall(
REDIS_KEYS['PLUGIN_SYS_VUL'].format(ident=ident))
if bplugin:
plugin = {'from_cache': True}
for k, v in bplugin.items():
plugin[k] = int(v) if k in ('id', 'level', 'status') else v
return cls(**plugin)
return None
def clean(self):
obj = self.get_by_key(self.name, 'name')
if self.id:
if int(self.id) != int(obj.id):
raise ValidatorException('名称已经存在')
elif obj:
raise ValidatorException('名称已经存在')
obj = self.get_by_key(self.ident, 'ident')
if self.id:
if int(self.id) != int(obj.id):
raise ValidatorException('标识已经存在')
elif obj:
raise ValidatorException('标识已经存在')
@classmethod
def create_or_replace(cls,
ident,
name,
level,
type,
params,
remark,
links,
id=None):
obj = None
if id:
obj = cls.get_by_key(id)
if obj is None:
obj = cls()
obj.ident = ident.strip()
obj.name = name.strip()
obj.level = level
obj.type = type
obj.params = json.dumps(params) if isinstance(params,
(dict, )) else params
obj.remark = remark.strip()
obj.links = links.strip()
obj.status = cls.STATUS_OK
has_error, errors = obj.valid()
if has_error:
return None, has_error, errors
db.session.add(obj)
db.session.commit()
obj.to_cache()
return obj, has_error, errors
@classmethod
def delete(cls, id):
obj = cls.get_by_key(id)
if obj:
obj.status = cls.STATUS_DELETE
db.session.add(obj)
db.session.commit()
obj.to_cache(True)
return obj
class PluginConfig(db.Model, ValidatorMixin, AsDictMixin):
STATUS_OK = 1
STATUS_DELETE = 2
id = db.Column(db.BigInteger, primary_key=True)
key = db.Column(db.String(64), nullable=False, default='')
value = db.Column(db.Text, nullable=False, default='')
status = db.Column(db.Integer, nullable=False)
created_time = db.Column(db.DateTime, nullable=False)
def __init__(self,
key,
value=None,
status=None,
created_time=None,
from_cache=False,
id=None):
self.id = id
self.key = key
self.value = value
self.status = self.STATUS_OK if status is None else status
self.created_time = timezone.now(
) if created_time is None else created_time
self.from_cache = from_cache
def clean(self):
obj = self.get_by_key(self.key, 'key')
if self.id:
if int(self.id) != int(obj.id):
raise ValidatorException('键已经存在')
elif obj:
raise ValidatorException('键已经存在')
@property
def value_object(self):
try:
return json.loads(self.value)
except BaseException as e:
return self.value
@classmethod
def all(cls):
return cls.query.filter(cls.status != cls.STATUS_DELETE).all()
@classmethod
def get_by_key(cls, value, key='id', all=False):
params = {key: value}
if not all:
params['status'] = cls.STATUS_OK
return cls.query.filter_by(**params).first()
@classmethod
def create_or_replace(cls, key, value, id=None):
obj = None
if id:
obj = cls.get_by_key(id)
if obj is None:
obj = cls(key=key)
obj.value = value.strip()
obj.status = cls.STATUS_OK
has_error, errors = obj.valid()
if has_error:
return None, has_error, errors
db.session.add(obj)
db.session.commit()
obj.to_cache()
return obj, has_error, errors
def to_cache(self, clear=False):
if clear:
redis.delete(REDIS_KEYS['PLUGIN_CONFIG'].format(ident=self.key))
else:
redis.hmset(REDIS_KEYS['PLUGIN_CONFIG'].format(ident=self.key),
self.as_dict_string())
return True
@classmethod
def get_by_cache(cls, key):
bconfig = redis.hgetall(REDIS_KEYS['PLUGIN_CONFIG'].format(ident=key))
if bconfig:
plugin = {'from_cache': True}
for k, v in bconfig.items():
plugin[k] = int(v) if k in ('id', 'status') else v
return cls(**plugin)
return None
@classmethod
def delete(cls, id):
obj = cls.get_by_key(id)
if obj:
obj.status = cls.STATUS_DELETE
db.session.add(obj)
db.session.commit()
obj.to_cache(True)
return obj
class Executor(db.Model, AsDictMixin):
STATUS_OK = 1
STATUS_DELETE = 2
id = db.Column(db.Integer, primary_key=True)
ident = db.Column(db.String(128), nullable=False, default='')
hostname = db.Column(db.String(128), nullable=False, default='')
pid = db.Column(db.Integer, nullable=False, default=0)
type = db.Column(db.Integer, nullable=False, default=0)
total = db.Column(db.Integer, nullable=False, default=0)
busy = db.Column(db.Integer, nullable=False, default=0)
idle = db.Column(db.Integer, nullable=False, default=0)
created_time = db.Column(db.DateTime)
heartbeat_time = db.Column(db.DateTime)
status = db.Column(db.Integer, nullable=False)
def __init__(self,
ident,
hostname,
pid,
type,
total=0,
busy=0,
idle=0,
running=0,
created_time=None,
heartbeat_time=None,
status=None):
self.ident = ident
self.hostname = hostname
self.pid = pid
self.type = type
self.total = total
self.busy = busy
self.idle = idle
self.created_time = timezone.now(
) if created_time is None else created_time
self.heartbeat_time = heartbeat_time
self.status = self.STATUS_OK if status is None else status
@classmethod
def stats(cls):
return {
'total':
db.session.query(cls).filter(cls.status == cls.STATUS_OK).count(),
'online':
db.session.query(cls).filter(
cls.status == cls.STATUS_OK, cls.heartbeat_time >=
timezone.now() - timedelta(minutes=3)).count()
}
@property
def inuse(self):
return max(self.busy, self.running)
@property
def type_text(self):
texts = {
2: '资产发现',
3: '漏洞扫描',
}
return texts.get(self.type)
@property
def running(self):
key = '{ident}:{type}'.format(ident=self.ident, type=self.type)
value = redis.hget(REDIS_KEYS['JOB_EXECUTOR_RUNNING'], key)
return 0 if value is None else int(value)
def running_incr(self):
key = '{ident}:{type}'.format(ident=self.ident, type=self.type)
return redis.hincrby(REDIS_KEYS['JOB_EXECUTOR_RUNNING'], key, 1)
def running_decr(self):
key = '{ident}:{type}'.format(ident=self.ident, type=self.type)
return redis.hincrby(REDIS_KEYS['JOB_EXECUTOR_RUNNING'], key, -1)
@classmethod
def running_decr_by_ident(cls, ident, type):
obj = cls.query.filter_by(ident=ident, type=type).first()
if obj:
obj.running_decr()
return True
@classmethod
def all(cls, online=True):
query = db.session.query(cls).filter(cls.status == cls.STATUS_OK)
if online:
query = query.filter(
cls.heartbeat_time >= timezone.now() - timedelta(minutes=3))
return query.all()
@classmethod
def create(cls, ident, hostname, pid, type, total, busy, idle):
executor = cls.query.filter_by(ident=ident, type=type).first()
if executor is None:
executor = cls(ident=ident,
hostname=hostname,
pid=pid,
type=type,
total=total,
busy=busy,
idle=idle)
executor.total = total
executor.busy = busy
executor.pid = pid
executor.heartbeat_time = timezone.now()
db.session.add(executor)
db.session.commit()
return executor
@classmethod
def get_by_key(cls, value, key='id'):
return db.session.query(cls).filter_by(**{key: value}).first()
def execute(self, sub_job):
self.running_incr()
key = REDIS_KEYS['JOB_EXECUTOR'].format(type=sub_job.type,
ident=self.ident)
redis.lpush(key, json.dumps(sub_job.as_dict()))
sub_job.start(self.ident)
return True
def as_dict(self):
rt = super(Executor, self).as_dict()
rt['running'] = self.running
return rt
@classmethod
def delete_by_key(cls, id):
executor = cls.get_by_key(id)
if executor:
executor.status = cls.STATUS_DELETE
db.session.add(executor)
db.session.commit()
return True
class Job(db.Model, AsDictMixin):
STATUS_WATING = 1
STATUS_PREPROCESS = 2
STATUS_DOING = 3
STATUS_CANCEL = 4
STATUS_SUCCESS = 5
STATUS_FAILURE = 6
STATUS_DELETE = 7
PROGRESS_PREPROCESS = 0
PROGRESS_DOING = 5
PROGRESS_COMPLATE = 100
id = db.Column(db.BigInteger, primary_key=True)
name = db.Column(db.String(64), nullable=False, default='')
user_id = db.Column(db.Integer, db.ForeignKey('user.id'))
type = db.Column(db.Integer, nullable=False)
job_params = db.Column(db.Text)
progress = db.Column(db.Integer, nullable=False, default=0)
status = db.Column(db.Integer, nullable=False)
explain = db.Column(db.String(512), nullable=False, default='')
created_time = db.Column(db.DateTime, nullable=False)
start_time = db.Column(db.DateTime)
finish_time = db.Column(db.DateTime)
sub_jobs = db.relationship('SubJob', backref="job", lazy='dynamic')
def __init__(self,
user=None,
type=None,
name='',
job_params=None,
status=None,
explain='',
created_time=None,
start_time=None,
finish_time=None,
progress=0,
id=None,
user_id=None,
from_cache=False,
*args,
**kwargs):
self.id = id
self.name = name
self.user_id = user.id if user else user_id
self.type = type
job_params = {} if job_params is None else job_params
self.job_params = json.dumps(job_params) if isinstance(
job_params, (dict, )) else job_params
self.progress = progress
self.status = self.STATUS_WATING if status is None else status
self.explain = explain
self.created_time = timezone.now(
) if created_time is None else created_time
self.start_time = start_time
self.finish_time = finish_time
self.from_cache = from_cache
@property
def job_params_object(self):
try:
return json.loads(self.job_params)
except BaseException as e:
return {}
@classmethod
def create(cls, user, type, name, job_params):
job = cls(user=user, type=type, name=name, job_params=job_params)
db.session.add(job)
db.session.commit()
redis.hmset(REDIS_KEYS['JOB_CONTENT'].format(id=job.id),
job.as_dict_string())
redis.lpush(REDIS_KEYS['JOB_PREPROCESS'], job.id)
return job
def preprocess(self):
job = self
if getattr(self, 'from_cache', False):
job = self.get_by_key(self.id)
job.status = self.STATUS_PREPROCESS
job.progress = self.PROGRESS_PREPROCESS
db.session.add(job)
db.session.commit()
redis.hmset(REDIS_KEYS['JOB_CONTENT'].format(id=self.id), {
'status': self.STATUS_PREPROCESS,
'progress': self.PROGRESS_PREPROCESS
})
redis.lpush(REDIS_KEYS['JOB_DOING'], self.id)
def start(self):
job = self
if getattr(self, 'from_cache', False):
job = self.get_by_key(self.id)
if job.start_time is None:
job.start_time = timezone.now()
job.status = self.STATUS_DOING
redis_content = {'status': self.STATUS_DOING}
if job.progress < self.PROGRESS_DOING:
job.progress = self.PROGRESS_DOING
redis_content['progress'] = self.PROGRESS_DOING
db.session.add(job)
db.session.commit()
redis.hmset(REDIS_KEYS['JOB_CONTENT'].format(id=self.id), {
'status': self.STATUS_DOING,
'progress': self.PROGRESS_DOING
})
@classmethod
def get_by_key(cls, value, key='id'):
return db.session.query(cls).filter_by(**{key: value}).first()
@classmethod
def get_by_cache(cls, id):
bjob = redis.hgetall(REDIS_KEYS['JOB_CONTENT'].format(id=id))
if bjob:
job = {'from_cache': True}
for k, v in bjob.items():
job[k] = int(v) if k in ('id', 'user_id', 'type', 'progress',
'status') else v
return cls(**job)
return None
def finish(self, status=None, explain=''):
job = self
if getattr(self, 'from_cache', False):
job = self.get_by_key(self.id)
job.finish_time = timezone.now()
job.status = self.STATUS_SUCCESS if status is None else self.STATUS_FAILURE
job.progress = self.PROGRESS_COMPLATE
job.explain = explain
db.session.add(job)
db.session.commit()
redis.delete(REDIS_KEYS['JOB_CONTENT'].format(id=self.id))
redis.lrem(REDIS_KEYS['JOB_DOING'], 0, self.id)
for _, key in JobType.KEYS.items():
redis.delete(REDIS_KEYS['JOB_QUEUE'].format(type=key, id=self.id))
return True
def cancel(self):
job = self
if getattr(self, 'from_cache', False):
job = self.get_by_key(self.id)
job.status = self.STATUS_CANCEL
job.progress = self.PROGRESS_COMPLATE
SubJob.cancel(job)
db.session.add(job)
db.session.commit()
redis.delete(REDIS_KEYS['JOB_CONTENT'].format(id=self.id))
redis.lrem(REDIS_KEYS['JOB_DOING'], 0, self.id)
for _, key in JobType.KEYS.items():
redis.delete(REDIS_KEYS['JOB_QUEUE'].format(type=key, id=self.id))
return True
def delete(self):
job = self
if getattr(self, 'from_cache', False):
job = self.get_by_key(self.id)
job.status = self.STATUS_DELETE
job.progress = self.PROGRESS_COMPLATE
SubJob.delete(job)
db.session.add(job)
db.session.commit()
redis.delete(REDIS_KEYS['JOB_CONTENT'].format(id=self.id))
redis.lrem(REDIS_KEYS['JOB_DOING'], 0, self.id)
for _, key in JobType.KEYS.items():
redis.delete(REDIS_KEYS['JOB_QUEUE'].format(type=key, id=self.id))
return True
def doing(self, progress=None):
job = self
if getattr(self, 'from_cache', False):
job = self.get_by_key(self.id)
job.status = self.STATUS_DOING
if progress is not None:
job.progress = progress
db.session.add(job)
db.session.commit()
class SubJob(db.Model, AsDictMixin):
STATUS_WATING = 1
STATUS_DOING = 2
STATUS_CANCEL = 3
STATUS_SUCCESS = 4
STATUS_FAILURE = 5
STATUS_DELETE = 6
id = db.Column(db.BigInteger, primary_key=True)
job_id = db.Column(db.BigInteger, db.ForeignKey('job.id'))
type = db.Column(db.Integer, nullable=False)
job_params = db.Column(db.Text)
created_time = db.Column(db.DateTime, nullable=False)
start_time = db.Column(db.DateTime)
finish_time = db.Column(db.DateTime)
status = db.Column(db.Integer, nullable=False)
executor = db.Column(db.String(128), nullable=False, default='')
explain = db.Column(db.String(512), nullable=False, default='')
def __init__(self,
job,
type,
job_params=None,
created_time=None,
start_time=None,
finish_time=None,
status=None,
executor='',
explain='',
*args,
**kwargs):
self.job_id = job.id
self.type = type
job_params = {} if job_params is None else job_params
self.job_params = json.dumps(job_params) if isinstance(
job_params, (dict, )) else job_params
self.created_time = timezone.now(
) if created_time is None else created_time
self.start_time = start_time
self.finish_time = finish_time
self.status = self.STATUS_WATING if status is None else status
self.executor = executor
self.explain = explain
@property
def job_params_object(self):
try:
return json.loads(self.job_params)
except BaseException as e:
return {}
@classmethod
def create(cls, job, type, job_params):
sub_job = cls(job=job, type=type, job_params=job_params)
db.session.add(sub_job)
db.session.commit()
redis.zadd(
REDIS_KEYS['JOB_QUEUE'].format(type=JobType.KEYS.get(type),
id=job.id), 0, sub_job.id)
return sub_job
@classmethod
def get_by_key(cls, value, key='id'):
return db.session.query(cls).filter_by(**{key: value}).first()
@classmethod
def cancel(cls, job):
for obj in cls.query.filter(
cls.job_id == job.id,
cls.status.notin_([
cls.STATUS_SUCCESS, cls.STATUS_CANCEL, cls.STATUS_FAILURE,
cls.STATUS_DELETE
])).all():
obj.status = cls.STATUS_CANCEL
Executor.running_decr_by_ident(obj.executor, obj.type)
db.session.add(obj)
db.session.commit()
return True
@classmethod
def delete(cls, job):
for obj in cls.query.filter(
cls.job_id == job.id,
cls.status.notin_([
cls.STATUS_SUCCESS, cls.STATUS_CANCEL, cls.STATUS_FAILURE,
cls.STATUS_DELETE
])).all():
obj.status = cls.STATUS_DELETE
Executor.running_decr_by_ident(obj.executor, obj.type)
db.session.add(obj)
db.session.commit()
return True
def start(self, executor=None):
self.status = self.STATUS_DOING
self.executor = executor if executor else self.executor
self.start_time = timezone.now()
redis.zadd(
REDIS_KEYS['JOB_QUEUE'].format(type=JobType.KEYS.get(self.type),
id=self.job.id), -1, self.id)
self.job.start()
db.session.add(self)
db.session.commit()
def finish(self, status=None, explain=''):
if self.status not in [
self.STATUS_SUCCESS, self.STATUS_CANCEL, self.STATUS_FAILURE,
self.STATUS_DELETE
]:
Executor.running_decr_by_ident(self.executor, self.type)
self.status = self.STATUS_SUCCESS if status is None else status
self.finish_time = timezone.now()
self.explain = explain
db.session.add(self)
db.session.commit()
redis.zadd(
REDIS_KEYS['JOB_QUEUE'].format(type=JobType.KEYS.get(self.type),
id=self.job.id), 1, self.id)
class User(db.Model, ValidatorMixin):
STATUS_REGISTED = 0
STATUS_LOCKED = 1
STATUS_OK = 2
STATUS_DELETE = 3
id = db.Column(db.Integer, primary_key=True)
name = db.Column(db.String(32), nullable=False, default='')
password = db.Column(db.String(512), nullable=False, default='')
email = db.Column(db.String(64), nullable=False, default='')
role_id = db.Column(db.Integer, db.ForeignKey('role.id'))
created_time = db.Column(db.DateTime, nullable=False)
last_login_time = db.Column(db.DateTime)
is_super = db.Column(db.Boolean, nullable=False, default=False)
status = db.Column(db.Integer, nullable=False)
def __init__(self, name, email, is_super=False, created_time=None, last_login_time=None, status=None, id=None, password=None, role_id=None):
self.id = id
self.name = name
self.password = password
self.email = email
self.role_id = role_id
self.created_time = created_time if created_time else timezone.now()
self.last_login_time = last_login_time
self.is_super = is_super
self.status = self.STATUS_REGISTED if status is None else status
@classmethod
def authenticate(cls, name, password):
user = cls.query.filter_by(name=name).first()
if not user:
user = cls.query.filter_by(email=name).first()
if user and check_password_hash(user.password, password):
user.last_login_time = timezone.now()
db.session.add(user)
db.session.commit()
return user
return None
def set_password(self, password):
self.password = generate_password_hash(password)
def clean_name(self):
name = self.name.strip()
if not ValidatorUtils.is_length(name, 3, 32):
raise ValidatorException('用户名必须为3到32个字符')
user = User.query.filter_by(name=self.name).first()
if user and (self.id is None or user.id != self.id):
raise ValidatorException('用户名已存在')
return name
def clean_email(self):
email = self.email.strip()
if not ValidatorUtils.is_email(email):
raise ValidatorException('邮箱格式不正确')
user = User.query.filter_by(email=self.email).first()
if user and (self.id is None or user.id != self.id):
raise ValidatorException('邮箱已存在')
return email
@classmethod
def create(cls, name, password, email, is_super):
user = cls(name=name, email=email, is_super=is_super, status=cls.STATUS_OK)
user.set_password(password)
has_error, errors = user.valid()
if has_error:
return None, has_error, errors
db.session.add(user)
db.session.commit()
return user, has_error, errors
@classmethod
def get_by_key(cls, value, key='id'):
return cls.query.filter_by(**{key:value}).first()
@classmethod
def all(cls):
return cls.query.filter(cls.status!=cls.STATUS_DELETE).all()