132 lines
4.8 KiB
Python
132 lines
4.8 KiB
Python
import random
|
|
import uuid
|
|
from datetime import datetime
|
|
|
|
from jinja2 import Template
|
|
|
|
from CTFd.utils import get_config
|
|
from CTFd.models import db
|
|
from CTFd.plugins.dynamic_challenges import DynamicChallenge
|
|
|
|
|
|
class WhaleConfig(db.Model):
|
|
key = db.Column(db.String(length=128), primary_key=True)
|
|
value = db.Column(db.Text)
|
|
|
|
def __init__(self, key, value):
|
|
self.key = key
|
|
self.value = value
|
|
|
|
def __repr__(self):
|
|
return "<WhaleConfig {0} {1}>".format(self.key, self.value)
|
|
|
|
|
|
class WhaleRedirectTemplate(db.Model):
|
|
key = db.Column(db.String(20), primary_key=True)
|
|
frp_template = db.Column(db.Text)
|
|
access_template = db.Column(db.Text)
|
|
|
|
def __init__(self, key, access_template, frp_template):
|
|
self.key = key
|
|
self.access_template = access_template
|
|
self.frp_template = frp_template
|
|
|
|
def __repr__(self):
|
|
return "<WhaleRedirectTemplate {0}>".format(self.key)
|
|
|
|
|
|
class DynamicDockerChallenge(DynamicChallenge):
|
|
__mapper_args__ = {"polymorphic_identity": "dynamic_docker"}
|
|
id = db.Column(
|
|
db.Integer, db.ForeignKey("dynamic_challenge.id", ondelete="CASCADE"), primary_key=True
|
|
)
|
|
|
|
memory_limit = db.Column(db.Text, default="128m")
|
|
cpu_limit = db.Column(db.Float, default=0.5)
|
|
dynamic_score = db.Column(db.Integer, default=0)
|
|
|
|
docker_image = db.Column(db.Text, default=0)
|
|
redirect_type = db.Column(db.Text, default=0)
|
|
redirect_port = db.Column(db.Integer, default=0)
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
kwargs["initial"] = kwargs["value"]
|
|
super(DynamicDockerChallenge, self).__init__(**kwargs)
|
|
|
|
|
|
class WhaleContainer(db.Model):
|
|
id = db.Column(db.Integer, primary_key=True, autoincrement=True)
|
|
user_id = db.Column(None, db.ForeignKey("users.id"))
|
|
challenge_id = db.Column(None, db.ForeignKey("challenges.id"))
|
|
start_time = db.Column(db.DateTime, nullable=False, default=datetime.utcnow)
|
|
renew_count = db.Column(db.Integer, nullable=False, default=0)
|
|
status = db.Column(db.Integer, default=1)
|
|
uuid = db.Column(db.String(256))
|
|
port = db.Column(db.Integer, nullable=True, default=0)
|
|
flag = db.Column(db.String(128), nullable=False)
|
|
|
|
# Relationships
|
|
user = db.relationship(
|
|
"Users", foreign_keys="WhaleContainer.user_id", lazy="select")
|
|
challenge = db.relationship(
|
|
"DynamicDockerChallenge", foreign_keys="WhaleContainer.challenge_id", lazy="select"
|
|
)
|
|
|
|
@property
|
|
def http_subdomain(self):
|
|
return Template(get_config(
|
|
'whale:template_http_subdomain', '{{ container.uuid }}'
|
|
)).render(container=self)
|
|
|
|
def __init__(self, user_id, challenge_id):
|
|
self.user_id = user_id
|
|
self.challenge_id = challenge_id
|
|
self.start_time = datetime.now()
|
|
self.renew_count = 0
|
|
self.uuid = str(uuid.uuid4())
|
|
self.flag = Template(get_config(
|
|
'whale:template_chall_flag', '{{ "flag{"+uuid.uuid4()|string+"}" }}'
|
|
)).render(container=self, uuid=uuid, random=random, get_config=get_config)
|
|
|
|
@property
|
|
def user_access(self):
|
|
return Template(WhaleRedirectTemplate.query.filter_by(
|
|
key=self.challenge.redirect_type
|
|
).first().access_template).render(container=self, get_config=get_config)
|
|
|
|
@property
|
|
def frp_config(self):
|
|
return Template(WhaleRedirectTemplate.query.filter_by(
|
|
key=self.challenge.redirect_type
|
|
).first().frp_template).render(container=self, get_config=get_config)
|
|
|
|
def __repr__(self):
|
|
return "<WhaleContainer ID:{0} {1} {2} {3} {4}>".format(self.id, self.user_id, self.challenge_id,
|
|
self.start_time, self.renew_count)
|
|
|
|
class WhaleContainerLog(db.Model):
|
|
id = db.Column(db.Integer, primary_key=True, autoincrement=True)
|
|
user_id = db.Column(None, db.ForeignKey("users.id"))
|
|
challenge_id = db.Column(None, db.ForeignKey("challenges.id"))
|
|
start_time = db.Column(db.DateTime, nullable=False, default=datetime.utcnow)
|
|
uuid = db.Column(db.String(256))
|
|
flag = db.Column(db.String(128), nullable=False)
|
|
|
|
# Relationships
|
|
user = db.relationship(
|
|
"Users", foreign_keys="WhaleContainerLog.user_id", lazy="select")
|
|
challenge = db.relationship(
|
|
"DynamicDockerChallenge", foreign_keys="WhaleContainerLog.challenge_id", lazy="select"
|
|
)
|
|
|
|
def __init__(self, container):
|
|
self.user_id = container.user_id
|
|
self.challenge_id = container.challenge_id
|
|
self.start_time = container.start_time
|
|
self.uuid = container.uuid
|
|
self.flag = container.flag
|
|
|
|
def __repr__(self):
|
|
return "<WhaleContainer ID:{0} {1} {2} {3} {4} {5}>".format(self.id, self.user_id, self.challenge_id,
|
|
self.start_time, self.uuid, self.flag)
|