"""JSON encoders to help with sending messages to server.""" import json from ipaddress import IPv4Address, IPv4Network, IPv6Address, IPv6Network from playhouse.shortcuts import model_to_dict from defence360agent.model import Model def ip_net_to_string(net) -> str: """ IPv4Network('192.168.1.1/32') -> '192.168.1.1' IPv4Network('192.168.1.0/24') -> '192.168.1.0/24' """ if not int(net.hostmask): return str(net.network_address) return str(net) class IPEncoder(json.JSONEncoder): def default(self, obj): if isinstance(obj, (IPv4Network, IPv6Network)): return ip_net_to_string(obj) if isinstance(obj, (IPv4Address, IPv6Address)): return str(obj) return json.JSONEncoder.default(self, obj) class ServerJSONEncoder(IPEncoder): def default(self, obj): if isinstance(obj, Model): return model_to_dict(obj) return super().default(obj)