From 10823e1c37abcb1276b67dab6214e847ed51818c Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <jstretch@digitalocean.com>
Date: Tue, 25 Apr 2017 13:00:28 -0400
Subject: [PATCH] Got rudimentary custom field creates/updates working

---
 netbox/dcim/api/serializers.py    |  4 +-
 netbox/extras/api/customfields.py | 79 +++++++++++++++++++++++++++----
 2 files changed, 72 insertions(+), 11 deletions(-)

diff --git a/netbox/dcim/api/serializers.py b/netbox/dcim/api/serializers.py
index a3149f963..de93f836b 100644
--- a/netbox/dcim/api/serializers.py
+++ b/netbox/dcim/api/serializers.py
@@ -66,13 +66,13 @@ class NestedSiteSerializer(serializers.ModelSerializer):
         fields = ['id', 'url', 'name', 'slug']
 
 
-class WritableSiteSerializer(serializers.ModelSerializer):
+class WritableSiteSerializer(CustomFieldModelSerializer):
 
     class Meta:
         model = Site
         fields = [
             'id', 'name', 'slug', 'region', 'tenant', 'facility', 'asn', 'physical_address', 'shipping_address',
-            'contact_name', 'contact_phone', 'contact_email', 'comments',
+            'contact_name', 'contact_phone', 'contact_email', 'comments', 'custom_fields',
         ]
 
 
diff --git a/netbox/extras/api/customfields.py b/netbox/extras/api/customfields.py
index e47eb41ab..4b96cd00c 100644
--- a/netbox/extras/api/customfields.py
+++ b/netbox/extras/api/customfields.py
@@ -1,6 +1,8 @@
 from django.contrib.contenttypes.models import ContentType
+from django.db import transaction
 
 from rest_framework import serializers
+from rest_framework.exceptions import ValidationError
 
 from extras.models import CF_TYPE_SELECT, CustomField, CustomFieldChoice, CustomFieldValue
 
@@ -14,6 +16,24 @@ class CustomFieldsSerializer(serializers.BaseSerializer):
     def to_representation(self, obj):
         return obj
 
+    def to_internal_value(self, data):
+
+        parent_content_type = ContentType.objects.get_for_model(self.parent.Meta.model)
+
+        for custom_field, value in data.items():
+
+            # Validate custom field name
+            try:
+                cf = CustomField.objects.get(name=custom_field)
+            except CustomField.DoesNotExist:
+                raise ValidationError(u"Unknown custom field: {}".format(custom_field))
+
+            # Validate custom field content type
+            if parent_content_type not in cf.obj_type.all():
+                raise ValidationError(u"Invalid custom field for {} objects".format(parent_content_type))
+
+        return data
+
 
 class CustomFieldModelSerializer(serializers.ModelSerializer):
     """
@@ -34,16 +54,57 @@ class CustomFieldModelSerializer(serializers.ModelSerializer):
 
         super(CustomFieldModelSerializer, self).__init__(*args, **kwargs)
 
-        # Retrieve the set of CustomFields which apply to this type of object
-        content_type = ContentType.objects.get_for_model(self.Meta.model)
-        fields = CustomField.objects.filter(obj_type=content_type)
+        if self.instance is not None:
 
-        # Populate CustomFieldValues for each instance from database
-        try:
-            for obj in self.instance:
-                _populate_custom_fields(obj, fields)
-        except TypeError:
-            _populate_custom_fields(self.instance, fields)
+            # Retrieve the set of CustomFields which apply to this type of object
+            content_type = ContentType.objects.get_for_model(self.Meta.model)
+            fields = CustomField.objects.filter(obj_type=content_type)
+
+            # Populate CustomFieldValues for each instance from database
+            try:
+                for obj in self.instance:
+                    _populate_custom_fields(obj, fields)
+            except TypeError:
+                _populate_custom_fields(self.instance, fields)
+
+    def _save_custom_fields(self, instance, custom_fields):
+        content_type = ContentType.objects.get_for_model(self.Meta.model)
+        for field_name, value in custom_fields.items():
+            custom_field = CustomField.objects.get(name=field_name)
+            CustomFieldValue.objects.update_or_create(
+                field=custom_field,
+                obj_type=content_type,
+                obj_id=instance.pk,
+                defaults={'serialized_value': value},
+            )
+
+    def create(self, validated_data):
+
+        custom_fields = validated_data.pop('custom_fields')
+
+        with transaction.atomic():
+
+            instance = super(CustomFieldModelSerializer, self).create(validated_data)
+
+            # Save custom fields
+            self._save_custom_fields(instance, custom_fields)
+            instance.custom_fields = custom_fields
+
+        return instance
+
+    def update(self, instance, validated_data):
+
+        custom_fields = validated_data.pop('custom_fields')
+
+        with transaction.atomic():
+
+            instance = super(CustomFieldModelSerializer, self).update(instance, validated_data)
+
+            # Save custom fields
+            self._save_custom_fields(instance, custom_fields)
+            instance.custom_fields = custom_fields
+
+        return instance
 
 
 class CustomFieldChoiceSerializer(serializers.ModelSerializer):