Files
NetBox-Export/netbox_export/services/codec.py
T

231 lines
8.5 KiB
Python

from __future__ import annotations
import base64
import datetime
import decimal
import uuid
from pathlib import PurePosixPath
from django.contrib.contenttypes.fields import GenericForeignKey
from django.db import models
SKIP_FIELD_NAMES = {
"id",
"created",
"last_updated",
"lft",
"rght",
"tree_id",
"level",
}
def object_key(obj) -> str:
return f"{obj._meta.label_lower}:{obj.pk}"
def model_label(model) -> str:
return model._meta.label_lower
def encode_scalar(value):
if value is None or isinstance(value, (bool, int, float, str)):
return value
if isinstance(value, decimal.Decimal):
return {"$type": "decimal", "value": str(value)}
if isinstance(value, uuid.UUID):
return {"$type": "uuid", "value": str(value)}
if isinstance(value, datetime.datetime):
return {"$type": "datetime", "value": value.isoformat()}
if isinstance(value, datetime.date):
return {"$type": "date", "value": value.isoformat()}
if isinstance(value, datetime.time):
return {"$type": "time", "value": value.isoformat()}
if isinstance(value, datetime.timedelta):
return {"$type": "duration", "value": value.total_seconds()}
if isinstance(value, bytes):
return {"$type": "bytes", "value": base64.b64encode(value).decode("ascii")}
if isinstance(value, (list, tuple)):
return [encode_scalar(item) for item in value]
if isinstance(value, dict):
return {str(key): encode_scalar(item) for key, item in value.items()}
return str(value)
def decode_scalar(value):
if isinstance(value, list):
return [decode_scalar(item) for item in value]
if not isinstance(value, dict) or "$type" not in value:
if isinstance(value, dict):
return {key: decode_scalar(item) for key, item in value.items()}
return value
kind = value["$type"]
raw = value.get("value")
decoders = {
"decimal": decimal.Decimal,
"uuid": uuid.UUID,
"datetime": datetime.datetime.fromisoformat,
"date": datetime.date.fromisoformat,
"time": datetime.time.fromisoformat,
"duration": lambda item: datetime.timedelta(seconds=item),
"bytes": lambda item: base64.b64decode(item.encode("ascii")),
}
return decoders[kind](raw)
def generic_foreign_keys(model) -> list[GenericForeignKey]:
return [field for field in model._meta.private_fields if isinstance(field, GenericForeignKey)]
def external_identity(obj) -> dict:
label = obj._meta.label_lower
if label in ("contenttypes.contenttype", "core.objecttype"):
return {"model": label, "lookup": {"app_label": obj.app_label, "model": obj.model}}
for field_name in ("username", "slug", "name"):
if hasattr(obj, field_name):
return {"model": label, "lookup": {field_name: encode_scalar(getattr(obj, field_name))}}
return {"model": label, "lookup": {"pk": encode_scalar(obj.pk)}}
def _reference_for_pk(model, pk, exported_keys: set[str]):
key = f"{model._meta.label_lower}:{pk}"
if key in exported_keys:
return {"ref": key}
try:
target = model._default_manager.get(pk=pk)
except model.DoesNotExist:
return None
return {"external": external_identity(target)}
def _encode_custom_field_data(obj, value: dict, exported_keys: set[str], custom_fields=None):
if custom_fields is None:
from extras.models import CustomField
custom_fields = CustomField.objects.get_for_model(type(obj))
definitions = {field.name: field for field in custom_fields}
encoded = {}
for name, raw_value in value.items():
custom_field = definitions.get(name)
if (
custom_field is None
or custom_field.type not in ("object", "multiobject")
or not custom_field.related_object_type
or raw_value in (None, "", [])
):
encoded[name] = encode_scalar(raw_value)
continue
target_model = custom_field.related_object_type.model_class()
if custom_field.type == "object":
reference = _reference_for_pk(target_model, raw_value, exported_keys)
encoded[name] = {"$type": "object_ref", "value": reference} if reference else None
else:
references = [
reference
for pk in raw_value
if (reference := _reference_for_pk(target_model, pk, exported_keys)) is not None
]
encoded[name] = {
"$type": "multiobject_ref",
"value": references,
}
return encoded
def _encode_custom_field_default(custom_field, value, exported_keys: set[str]):
if (
custom_field.type not in ("object", "multiobject")
or not custom_field.related_object_type
or value in (None, "", [])
):
return encode_scalar(value)
target_model = custom_field.related_object_type.model_class()
if custom_field.type == "object":
reference = _reference_for_pk(target_model, value, exported_keys)
return {"$type": "object_ref", "value": reference} if reference else None
references = [
reference
for pk in value
if (reference := _reference_for_pk(target_model, pk, exported_keys)) is not None
]
return {
"$type": "multiobject_ref",
"value": references,
}
def serialize_object(obj, exported_keys: set[str], assets: dict[str, bytes], custom_fields=None) -> dict:
record_id = object_key(obj)
gfk_fields = generic_foreign_keys(type(obj))
gfk_storage = {name for field in gfk_fields for name in (field.ct_field, field.fk_field)}
scalars = {}
relations = {}
files = {}
for field in obj._meta.concrete_fields:
if field.primary_key or field.name.startswith("_") or field.name in SKIP_FIELD_NAMES or field.name in gfk_storage:
continue
if isinstance(field, (models.ForeignKey, models.OneToOneField)):
related = getattr(obj, field.name, None)
if related is None:
relations[field.name] = None
elif object_key(related) in exported_keys:
relations[field.name] = {"ref": object_key(related)}
else:
relations[field.name] = {"external": external_identity(related)}
continue
if isinstance(field, models.FileField):
file_value = getattr(obj, field.name)
if not file_value:
files[field.name] = None
continue
filename = PurePosixPath(str(file_value.name)).name
archive_path = f"assets/{record_id.replace(':', '_')}/{field.name}/{filename}"
try:
with file_value.open("rb") as source:
assets[archive_path] = source.read()
except (FileNotFoundError, OSError, ValueError):
files[field.name] = {"missing": str(file_value.name)}
else:
files[field.name] = {"path": archive_path, "name": str(file_value.name)}
continue
value = field.value_from_object(obj)
if field.name == "custom_field_data" and isinstance(value, dict):
scalars[field.name] = _encode_custom_field_data(obj, value, exported_keys, custom_fields)
elif obj._meta.label_lower == "extras.customfield" and field.name == "default":
scalars[field.name] = _encode_custom_field_default(obj, value, exported_keys)
else:
scalars[field.name] = encode_scalar(value)
generic_relations = {}
for field in gfk_fields:
related = getattr(obj, field.name, None)
if related is None:
generic_relations[field.name] = None
elif object_key(related) in exported_keys:
generic_relations[field.name] = {"ref": object_key(related)}
else:
generic_relations[field.name] = {"external": external_identity(related)}
many_to_many = {}
for field in obj._meta.many_to_many:
try:
values = list(getattr(obj, field.name).all())
except (AttributeError, TypeError):
continue
many_to_many[field.name] = [
{"ref": object_key(item)} if object_key(item) in exported_keys else {"external": external_identity(item)}
for item in values
]
return {
"id": record_id,
"model": obj._meta.label_lower,
"source_pk": str(obj.pk),
"fields": scalars,
"relations": relations,
"generic_relations": generic_relations,
"many_to_many": many_to_many,
"files": files,
}