class AgentType(db.Model): __tablename__ = "AgentType" Id = db.Column(db.Integer, primary_key=True) Agents = db.relationship('Agent', backref='AgentType', lazy=True) Payloads = db.relationship('Payload', backref='AgentType', lazy=True) AgentTransportTypes = db.relationship('AgentTransportType', backref='AgentType', lazy=True) AgentTypeArchitectures = db.relationship('AgentTypeArchitecture', backref='AgentType', lazy=True) AgentTypeConfigurations = db.relationship('AgentTypeConfiguration', backref='AgentType', lazy=True) AgentTypeFormats = db.relationship('AgentTypeFormat', backref='AgentType', lazy=True) AgentTypeOperatingSystems = db.relationship('AgentTypeOperatingSystem', backref='AgentType', lazy=True) AgentTypeVersions = db.relationship('AgentTypeVersion', backref='AgentType', lazy=True) Payloads = db.relationship('Payload', backref='AgentType', lazy=True) Commands = db.relationship('Command', backref='AgentType', lazy=True) Name = db.Column(db.String) Guid = db.Column(db.String) def __repr__(self): return '<AgentType: %s>' % str(self.Id)
class SetupModel(db.Model): __tablename__ = 'setup' setup_id = db.Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, unique=True, nullable=False) training_id = db.Column(UUID(as_uuid=True), db.ForeignKey('training.training_id'), nullable=False) bike_id = db.Column(UUID(as_uuid=True), db.ForeignKey('bike.bike_id'), nullable=False) operating_hours = db.Column(db.Float, nullable=False) weather_current = db.Column(JSON, nullable=True) slick_pressure_front = db.Column(db.Float, nullable=True) slick_pressure_rear = db.Column(db.Float, nullable=True) rain_pressure_front = db.Column(db.Float, nullable=True) rain_pressure_rear = db.Column(db.Float, nullable=True) setup = db.Column(JSON, nullable=True) comment = db.Column(db.Text, nullable=True) datetime_created = db.Column(db.DateTime, nullable=False, default=db.utcnow) datetime_last_modified = db.Column(db.DateTime, nullable=False, default=db.utcnow, onupdate=db.utcnow) datetime_display = db.Column(db.DateTime, nullable=False, default=db.utcnow) bike = db.relationship('BikeModel', backref=db.backref( 'setups', order_by='SetupModel.datetime_display.asc()')) training = db.relationship( 'TrainingModel', backref=db.backref('setups', order_by='SetupModel.datetime_display.asc()')) def __repr__(self): return f"Setup[" \ f"'{self.setup_id}': (" \ f"'{self.operating_hours}', " \ f"'{self.comment}', " \ f"'{self.datetime_display}', " \ f"training_id: '{self.training_id}', " \ f"bike_id: '{self.bike_id}'" \ f")]"
class SessionModel(db.Model): __tablename__ = 'session' session_id = db.Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, unique=True, nullable=False) training_id = db.Column(UUID(as_uuid=True), db.ForeignKey('training.training_id'), nullable=False) setup_id = db.Column(UUID(as_uuid=True), db.ForeignKey('setup.setup_id'), nullable=True) bike_id = db.Column(UUID(as_uuid=True), db.ForeignKey('bike.bike_id'), nullable=True) application = db.Column(db.String(50), nullable=True) datetime_created = db.Column(db.DateTime, nullable=False, default=db.utcnow) datetime_last_modified = db.Column(db.DateTime, nullable=False, default=db.utcnow, onupdate=db.utcnow) datetime_display = db.Column(db.DateTime, nullable=False, default=db.utcnow) bike = db.relationship('BikeModel', backref=db.backref( 'sessions', order_by='SessionModel.datetime_display.asc()')) training = db.relationship( 'TrainingModel', backref=db.backref('sessions', order_by='SessionModel.datetime_display.asc()')) setup = db.relationship( 'SetupModel', backref=db.backref('sessions', order_by='SessionModel.datetime_display.asc()')) def __repr__(self): return f"Session[" \ f"'{self.session_id}': (" \ f"'{self.datetime_display}', " \ f"training_id: '{self.training_id}', " \ f"bike_id: '{self.bike_id}', " \ f"setup_id: '{self.setup_id}', " \ f")]"
class Payload(db.Model): __tablename__ = "Payload" Id = db.Column(db.Integer, primary_key=True) Name = db.Column(db.String, unique=True) Description = db.Column(db.String) Key = db.Column(db.String) Built = db.Column(db.Boolean) BuildToken = db.Column(db.String) Filename = db.Column(db.String) AgentTypeId = db.Column(db.Integer, db.ForeignKey('AgentType.Id'), nullable=False) AgentTypeArchitectureId = db.Column( db.Integer, db.ForeignKey('AgentTypeArchitecture.Id'), nullable=False) AgentTypeConfigurationId = db.Column( db.Integer, db.ForeignKey('AgentTypeConfiguration.Id'), nullable=False) AgentTypeOperatingSystemId = db.Column( db.Integer, db.ForeignKey('AgentTypeOperatingSystem.Id'), nullable=False) AgentTypeVersionId = db.Column(db.Integer, db.ForeignKey('AgentTypeVersion.Id'), nullable=False) AgentTypeFormatId = db.Column(db.Integer, db.ForeignKey('AgentTypeFormat.Id'), nullable=False) AgentTransportTypeId = db.Column(db.Integer, db.ForeignKey('AgentTransportType.Id'), nullable=False) TransportId = db.Column(db.Integer, db.ForeignKey('Transport.Id'), nullable=False) BeaconInterval = db.Column(db.Integer) Jitter = db.Column(db.Float) Created = db.Column(db.DateTime) LastDownloaded = db.Column(db.DateTime) ExpirationDate = db.Column(db.DateTime) Enabled = db.Column(db.Boolean) Agents = db.relationship('Agent', backref='Payload', lazy=True) StagingMessages = db.relationship("StagingMessage", backref='StagingConfig', lazy=True) Visible = db.Column(db.Boolean) def __repr__(self): return '<Payload: %s>' % str(self.Id)
class Agent(db.Model): __tablename__ = "Agent" Id = db.Column(db.Integer, primary_key=True) Name = db.Column(db.String) StagingId = db.Column(db.String) AesPassword = db.Column(db.String) Username = db.Column(db.String) Hostname = db.Column(db.String) PID = db.Column(db.Integer) OperatingSystem = db.Column(db.String) Admin = db.Column(db.Boolean) AgentTypeId = db.Column(db.Integer, db.ForeignKey('AgentType.Id'), nullable=False) StagingResponseId = db.Column(db.Integer, db.ForeignKey('StagingMessage.Id'), nullable=False) PayloadId = db.Column(db.Integer, db.ForeignKey('Payload.Id'), nullable=False) TransportId = db.Column(db.Integer, db.ForeignKey('Transport.Id'), nullable=False) InternalIP = db.Column(db.String) ExternalIP = db.Column(db.String) InitialCheckin = db.Column(db.DateTime) LastCheckin = db.Column(db.DateTime) BeaconInterval = db.Column(db.Integer) Jitter = db.Column(db.Float) Tasks = db.relationship('AgentTask', backref='Agent', lazy=True) ConsoleMessages = db.relationship("ConsoleMessage", backref='Agent', lazy=True) AvailableTransports = db.relationship('Transport', secondary=AgentsTransportsXREF, lazy='subquery', backref=db.backref('AvailableAgents', lazy=True)) AvailableModules = db.relationship('Module', secondary=AgentModulesXREF, lazy='subquery', backref=db.backref('AvailableAgents', lazy=True)) Visible = db.Column(db.Boolean) def __repr__(self): if self.Name: return '<Agent: %s>' % self.Name else: return '<Agent: %s>' % str(self.Id)
class UserRole(db.Model): __tablename__ = "UserRole" Id = db.Column(db.Integer, primary_key=True) Name = db.Column(db.String, nullable=False, unique=True) Users = db.relationship('User', backref='UserRole', lazy=True) def __repr__(self): return '<Role: %s>' % str(self.Id)
class HistoryModel(db.Model): __tablename__ = 'history' history_id = db.Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, unique=True, nullable=False) maintenance_id = db.Column(UUID(as_uuid=True), db.ForeignKey('maintenance.maintenance_id'), nullable=False) bike_id = db.Column(UUID(as_uuid=True), db.ForeignKey('bike.bike_id'), nullable=False) operating_hours = db.Column(db.Float, nullable=False) comment = db.Column(db.Text, nullable=True) tags = db.Column(ARRAY(db.String), nullable=True) datetime_created = db.Column(db.DateTime, nullable=False, default=db.utcnow) datetime_last_modified = db.Column(db.DateTime, nullable=False, default=db.utcnow, onupdate=db.utcnow) datetime_display = db.Column(db.DateTime, nullable=False, default=db.utcnow) bike = db.relationship('BikeModel', backref=db.backref('history')) maintenance = db.relationship( 'MaintenanceModel', backref=db.backref('history', order_by='HistoryModel.operating_hours.desc()')) def __repr__(self): return f"History[" \ f"'{self.history_id}': (" \ f"'{self.operating_hours}', " \ f"'{self.comment}', " \ f"'{self.datetime_display}', " \ f"mtn_id: '{self.maintenance_id}'," \ f"bike_id: '{self.bike_id}'" \ f")]"
class Transport(db.Model): __tablename__ = "Transport" Id = db.Column(db.Integer, primary_key=True) Name = db.Column(db.String) TransportType = db.Column(db.String) Guid = db.Column(db.String) Created = db.Column(db.DateTime) LastCheckin = db.Column(db.DateTime) Configuration = db.Column(db.String) ApiKeyId = db.Column(db.Integer, db.ForeignKey('ApiKey.Id'), nullable=False) Enabled = db.Column(db.Boolean) Visible = db.Column(db.Boolean) Agents = db.relationship('Agent', backref='Transport', lazy=True) Payloads = db.relationship('Payload', backref='Transport', lazy=True) def __repr__(self): return '<Transport: %s>' % str(self.Id)
class Command(db.Model): __tablename__ = "Command" Id = db.Column(db.Integer, primary_key=True) Name = db.Column(db.String) Description = db.Column(db.String) Help = db.Column(db.String) MitreReference = db.Column(db.String) OpsecSafe = db.Column(db.Boolean) ModuleId = db.Column(db.Integer, db.ForeignKey('Module.Id')) AgentTypeId = db.Column(db.Integer, db.ForeignKey('AgentType.Id')) Parameters = db.relationship('CommandParameter', backref='Command', lazy=True)
class Session(db.Model): __tablename__ = 'session' id = db.Column(UUIDType(), primary_key=True) uid = db.Column(UUIDType(), db.ForeignKey('user.id'), nullable=False) user = db.relationship('User', backref=db.backref('sessions', lazy=True)) expire = db.Column(db.DateTime(), nullable=False) @staticmethod def new_session(username, expire=None) -> 'Session': if expire == None: expire = datetime.now() + timedelta(hours=3) if type(expire) in (int, float): expire = datetime.fromtimestamp(expire) user = User.get_by_username(username) return Session( id=uuid4(), user=user, expire=expire, ) @staticmethod def get_by_id(id) -> 'Session': if type(id) != UUID: try: id = UUID(id) except: raise AttributeError("Invalid UUID") query = Session.query.filter_by(id=id) if query.count() == 0: raise ValueError(f"Session {str(id)} not found.") session = query.first() if session.expired(): db.session.delete(session) db.session.commit() raise ValueError(f"Session {str(id)} not found.") return session def json(self) -> dict: return { "id": str(self.id), "user": self.user.json(), "expire": int(self.expire.timestamp()), } def expired(self) -> bool: return self.expire < datetime.now()
class AgentTypeOperatingSystem(db.Model): __tablename__ = "AgentTypeOperatingSystem" Id = db.Column(db.Integer, primary_key=True) Name = db.Column(db.String) AgentTypeId = db.Column(db.Integer, db.ForeignKey('AgentType.Id'), nullable=False) Payloads = db.relationship('Payload', backref='AgentTypeOperatingSystem', lazy=True) def __repr__(self): if self.Name: return '<AgentTypeOperatingSystem: %s>' % self.Name else: return '<AgentTypeOperatingSystem: %s>' % str(self.Id)
class AgentTaskUpdate(db.Model): __tablename__ = "AgentTaskUpdate" Id = db.Column(db.Integer, primary_key=True) AgentId = db.Column(db.Integer, db.ForeignKey('Agent.Id')) TaskId = db.Column(db.Integer, db.ForeignKey('AgentTask.Id')) Message = db.Column(db.String) Complete = db.Column(db.Boolean) Success = db.Column(db.Boolean) Received = db.Column(db.DateTime) IOCs = db.relationship("IOC", backref='AgentTaskUpdate', lazy=True) def __repr__(self): return '<AgentTaskUpdate: %s>' % str(self.Id)
class AgentTransportType(db.Model): __tablename__ = "AgentTransportType" Id = db.Column(db.Integer, primary_key=True) Name = db.Column(db.String) TransportTypeGuid = db.Column(db.String) BuildCommand = db.Column(db.String) BuildLocation = db.Column(db.String) AgentTypeId = db.Column(db.Integer, db.ForeignKey('AgentType.Id'), nullable=False) Payloads = db.relationship('Payload', backref='AgentTransportType', lazy=True) def __repr__(self): return '<AgentType: %s>' % str(self.Id)
class AgentTask(db.Model): __tablename__ = "AgentTask" Id = db.Column(db.Integer, primary_key=True) Name = db.Column(db.String) AgentId = db.Column(db.Integer, db.ForeignKey('Agent.Id'), nullable=False) ConsoleMessageId = db.Column(db.Integer, db.ForeignKey('ConsoleMessage.Id'), nullable=True) AgentTaskUpdates = db.relationship('AgentTaskUpdate', backref='AgentTask', lazy=True) Action = db.Column(db.String) Command = db.Column(db.String) Created = db.Column(db.DateTime) def __repr__(self): return '<Task: %s>' % str(self.Id)
class LaptimeModel(db.Model): __tablename__ = 'laptime' lap_id = db.Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, unique=True, nullable=False) session_id = db.Column(UUID(as_uuid=True), db.ForeignKey('session.session_id'), nullable=False) lap_no = db.Column(db.Integer, nullable=False) valid = db.Column(db.Boolean, nullable=False, default=True) track_layout = db.Column(db.String, nullable=False, default="A") laptime_seconds = db.Column(db.Float, nullable=False) sectors = db.Column(JSON, nullable=False) datetime_created = db.Column(db.DateTime, nullable=False, default=db.utcnow) datetime_last_modified = db.Column(db.DateTime, nullable=False, default=db.utcnow, onupdate=db.utcnow) datetime_display = db.Column(db.DateTime, nullable=False, default=db.utcnow) session = db.relationship( 'SessionModel', backref=db.backref('laptimes', order_by='LaptimeModel.datetime_display.asc()')) def __repr__(self): return f"Lap time[" \ f"'{self.lap_id}': (" \ f"'{self.datetime_display}', " \ f"training_id: '{self.training_id}', " \ f"bike_id: '{self.bike_id}', " \ f"setup_id: '{self.setup_id}', " \ f")]"
class SparepartModel(db.Model): __tablename__ = 'sparepart' sparepart_id = db.Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, unique=True, nullable=False) name = db.Column(db.String, nullable=False) items = db.relationship('SparepartitemModel', backref=db.backref('sparepart')) module = db.Column(db.String, nullable=False) min_stock = db.Column(db.Integer, nullable=True) current_stock = db.column_property( db.select([func.sum(SparepartitemModel.stock) ]).where(SparepartitemModel.sparepart_id == sparepart_id).correlate_except(SparepartitemModel)) datetime_created = db.Column(db.DateTime, nullable=False, default=db.utcnow) datetime_last_modified = db.Column(db.DateTime, nullable=False, default=db.utcnow, onupdate=db.utcnow) datetime_display = db.Column(db.DateTime, nullable=False, default=db.utcnow) def __repr__(self): return f"Sparepart[" \ f"'{self.sparepart_id}': (" \ f"'{self.module}', " \ f"'{self.alarm}', " \ f"'{self.datetime_display}'" \ f")]"
class User(db.Model): __tablename__ = "User" Id = db.Column(db.Integer, primary_key=True) Username = db.Column(db.String, unique=True) Password = db.Column(db.LargeBinary) RoleId = db.Column(db.Integer, db.ForeignKey('UserRole.Id'), nullable=False) ApiKeys = db.relationship("ApiKey", backref='User', lazy=True) Authenticated = db.Column(db.Boolean, default=False) ConsoleMessages = db.relationship("ConsoleMessage", backref='User', lazy=True) Files = db.relationship("FactionFile", backref='User', lazy=True) Created = db.Column(db.DateTime) LastLogin = db.Column(db.DateTime) Enabled = db.Column(db.Boolean) Visible = db.Column(db.Boolean) def is_active(self): """True, as all users are active.""" return True def get_id(self): """Return the email address to satisfy Flask-Login's requirements.""" return self.Username def is_authenticated(self): """Return True if the user is authenticated.""" return self.Authenticated def is_anonymous(self): """False, as anonymous users aren't supported.""" return False def change_password(self, current_password, new_password): log("change_password", "Got password change request") if bcrypt.checkpw(current_password.encode('utf-8'), self.Password) and self.Enabled: self.Password = bcrypt.hashpw(new_password.encode('utf-8'), bcrypt.gensalt()) db.session.add(self) db.session.commit() log("change_password", "Password changed") return dict({ "Success": True, "Message": 'Changed password for user: {0}'.format(self.Username) }) log("change_password", "Current password incorrect") return { 'Success':False, 'Message':'Invalid username or password.' } def get_api_keys(self): api_keys = self.ApiKeys log("get_api_keys", "Got api keys: {0}".format(str(api_keys))) results = [] for api_key in api_keys: result = dict() result['Id'] = api_key.Id result['Name'] = api_key.Name result['Created'] = None result['LastUsed'] = None result['Type'] = api_key.Type if api_key.Created: result['Created'] = api_key.Created.isoformat() if api_key.LastUsed: result['LastUsed'] = api_key.LastUsed.isoformat() results.append(result) return { 'Success':True, 'Results':results } def delete_api_key(self, api_key_id): api_keys = self.ApiKeys for api_key in api_keys: if api_key.Id == api_key_id: db.session.delete(api_key) db.session.commit() return { 'Success':True, 'Message':"Api Key {0} deleted".format(api_key.Name) } return { 'Success':False, 'Message':"Api Key ID: {0} not found".format(api_key_id) }