diff --git a/xsense/async_xsense.py b/xsense/async_xsense.py
index 210c24e..527db27 100644
--- a/xsense/async_xsense.py
+++ b/xsense/async_xsense.py
@@ -246,22 +246,9 @@ async def get_state(self, station: Station):
             raise APIFailure(f'Unable to retrieve station data: {self._lastres.status}/{text}')
 
     async def set_state(self, entity: Entity, shadow: str, topic: str, definition: Dict):
-        station = entity.station
-        t = datetime.now()
-        timestamp = t.strftime('%Y%m%d%H%M%S')
-
-        desired = {
-            "deviceSN": entity.sn,
-            "shadow": shadow,
-            "stationSN": station.sn,
-            "time": timestamp,
-            "userId": self.userid
-        }
-        desired.update(definition.get('extra', {}))
-
-        data = {"state": {"desired": desired}}
+        station, data = self.build_desired_state(entity, shadow, definition)
 
-        res = await self.do_thing(station, topic, data)
+        return await self.do_thing(station, topic, data)
 
     async def action(self, entity: Entity, action: str):
         entity_def = entities.get(entity.type)
@@ -275,4 +262,4 @@ async def action(self, entity: Entity, action: str):
         topic = action_def.get('topic')
         if callable(topic):
             topic = topic(entity)
-        await self.set_state(entity, action_def['shadow'], topic, action_def)
+        return await self.set_state(entity, action_def['shadow'], topic, action_def)
diff --git a/xsense/base.py b/xsense/base.py
index 007a59e..f5bbbb7 100644
--- a/xsense/base.py
+++ b/xsense/base.py
@@ -11,7 +11,7 @@
 
 from .entity import Entity
 from .entity_map import entities
-from .exceptions import AuthFailed
+from .exceptions import AuthFailed, XSenseError
 from .station import Station
 from .house import House
 
@@ -233,4 +233,30 @@ def _parse_get_house_state(self, house: House, data: Dict):
     def has_action(self, entity: Entity, action: str):
         if entity_def := entities.get(entity.type):
             return any(a for a in entity_def.get('actions', []) if a.get('action') == action)
-        return False
\ No newline at end of file
+        return False
+
+    def _station_for_entity(self, entity: Entity) -> Station:
+        station = getattr(entity, 'station', None)
+        if station:
+            return station
+        if isinstance(entity, Station):
+            return entity
+        raise XSenseError(f'Entity type {entity.type} has no station')
+
+    def build_desired_state(self, entity: Entity, shadow: str, definition: Dict):
+        station = self._station_for_entity(entity)
+        t = datetime.now(timezone.utc)
+        timestamp = t.strftime('%Y%m%d%H%M%S')
+
+        desired = {
+            "shadow": shadow,
+            "stationSN": station.sn,
+            "time": timestamp,
+            "userId": self.userid
+        }
+        if not isinstance(entity, Station):
+            desired["deviceSN"] = entity.sn
+        desired.update(definition.get('extra', {}))
+        desired.update(definition.get('data', {}))
+
+        return station, {"state": {"desired": desired}}
diff --git a/xsense/station.py b/xsense/station.py
index 8de3842..222657d 100644
--- a/xsense/station.py
+++ b/xsense/station.py
@@ -62,6 +62,9 @@ def set_alarm_data(self, values: dict):
             if v := values.get(k):
                 self._alarm_data[k] = v
 
+        if safe_mode := values.get('safeMode'):
+            self.safe_mode = safe_mode
+
     @property
     def alarm_data(self):
         return self._alarm_data
diff --git a/xsense/xsense.py b/xsense/xsense.py
index 8c239d1..81a96e0 100644
--- a/xsense/xsense.py
+++ b/xsense/xsense.py
@@ -1,4 +1,4 @@
-from datetime import datetime, timedelta
+from datetime import datetime
 from typing import Dict
 
 import requests
@@ -222,22 +222,9 @@ def get_state(self, station: Station):
             raise APIFailure(f'Unable to retrieve station data: {self._lastres.status_code}/{self._lastres.text}')
 
     def set_state(self, entity: Entity, shadow: str, topic: str, definition: Dict):
-        station = entity.station
-        t = datetime.now()
-        timestamp = t.strftime('%Y%m%d%H%M%S')
-
-        desired = {
-            "deviceSN": entity.sn,
-            "shadow": shadow,
-            "stationSN": station.sn,
-            "time": timestamp,
-            "userId": self.userid
-        }
-        desired.update(definition.get('extra', {}))
-
-        data = {"state": {"desired": desired}}
+        station, data = self.build_desired_state(entity, shadow, definition)
 
-        res = self.do_thing(station, topic, data)
+        return self.do_thing(station, topic, data)
 
     def action(self, entity: Entity, action: str):
         entity_def = entities.get(entity.type)
@@ -251,4 +238,4 @@ def action(self, entity: Entity, action: str):
         topic = action_def.get('topic')
         if callable(topic):
             topic = topic(entity)
-        self.set_state(entity, action_def['shadow'], topic, action_def)
+        return self.set_state(entity, action_def['shadow'], topic, action_def)