Add hostname validation to database model
This commit is contained in:
@@ -4,6 +4,7 @@ import logging
|
||||
import os
|
||||
|
||||
from . import utc_now
|
||||
from .dns import encode_dnsname, EncodingError
|
||||
from peewee import (
|
||||
AutoField,
|
||||
CharField,
|
||||
@@ -84,6 +85,14 @@ class Hostname(BaseModel):
|
||||
(('hostname', 'zone'), True),
|
||||
)
|
||||
|
||||
def save(self, *args, **kwargs):
|
||||
"""Validate and encode hostname/zone before saving."""
|
||||
if self.hostname:
|
||||
self.hostname = encode_dnsname(self.hostname)
|
||||
if self.zone:
|
||||
self.zone = encode_dnsname(self.zone)
|
||||
return super().save(*args, **kwargs)
|
||||
|
||||
|
||||
class Version(BaseModel):
|
||||
"""Database schema version for migrations."""
|
||||
@@ -301,7 +310,7 @@ def get_hostname_for_user(hostname: str, user: User):
|
||||
return Hostname.get((fqdn == hostname) & (Hostname.user == user))
|
||||
|
||||
|
||||
# Re-export DoesNotExist for convenience
|
||||
# Re-export DoesNotExist and EncodingError for convenience
|
||||
__all__ = [
|
||||
'db',
|
||||
'DATABASE_VERSION',
|
||||
@@ -314,4 +323,5 @@ __all__ = [
|
||||
'get_hostname',
|
||||
'get_hostname_for_user',
|
||||
'DoesNotExist',
|
||||
'EncodingError',
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user