diff --git a/disco/bot/__init__.py b/disco/bot/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/disco/gateway/client.py b/disco/gateway/client.py index 7398033..dc65919 100644 --- a/disco/gateway/client.py +++ b/disco/gateway/client.py @@ -7,10 +7,11 @@ from holster.emitter import Emitter # from holster.util import SimpleObject from disco.gateway.packets import OPCode, HeartbeatPacket, ResumePacket, IdentifyPacket -from disco.gateway.events import GatewayEvent +from disco.gateway.events import GatewayEvent, Ready from disco.util.logging import LoggingClass GATEWAY_VERSION = 6 +TEN_MEGABYTES = 10490000 def log_error(log, msg, w): @@ -29,12 +30,15 @@ class GatewayClient(LoggingClass): self.client = client self.emitter = Emitter(gevent.spawn) + self.emitter.on(Ready, self.on_ready) + # Websocket connection self.ws = None # State self.seq = 0 self.session_id = None + self.reconnects = 0 # Cached gateway URL self._cached_gateway_url = None @@ -56,8 +60,9 @@ class GatewayClient(LoggingClass): gevent.sleep(interval / 1000) def handle_dispatch(self, packet): - obj = GatewayEvent.from_dispatch(packet) - self.log.info('Got dispatch for %s', obj.user.id) + cls, obj = GatewayEvent.from_dispatch(packet) + self.log.info('Dispatching %s', cls) + self.emitter.emit(cls, obj) def handle_heartbeat(self, packet): pass @@ -75,6 +80,11 @@ class GatewayClient(LoggingClass): def handle_heartbeat_ack(self, packet): pass + def on_ready(self, ready): + self.log.info('Recieved READY') + self.session_id = ready.session_id + self.reconnects = 0 + def connect(self): if not self._cached_gateway_url: self._cached_gateway_url = self.client.api.gateway(version=GATEWAY_VERSION, encoding='json') @@ -89,10 +99,9 @@ class GatewayClient(LoggingClass): ) def on_message(self, ws, msg): - # Check if we're JSON + # Detect zlib and decompress if msg[0] != '{': - print 'zlib' - msg = zlib.decompress(msg) + msg = zlib.decompress(msg, 15, TEN_MEGABYTES) try: data = json.loads(msg) diff --git a/disco/gateway/events.py b/disco/gateway/events.py index 6a4ce71..8b6c079 100644 --- a/disco/gateway/events.py +++ b/disco/gateway/events.py @@ -1,8 +1,7 @@ import inflection import skema -from disco.types.user import User -from disco.types.guild import Guild +from disco.types import Guild, Channel, User, GuildMember, Role, Message, VoiceState class GatewayEvent(skema.Model): @@ -12,7 +11,25 @@ class GatewayEvent(skema.Model): if not cls: raise Exception('Could not find cls for {}'.format(obj['t'])) - return cls(obj['d']) + return cls, cls.create(obj['d']) + + @classmethod + def create(cls, obj): + self = cls(obj) + self.validate() + return self + + +def Sub(field): + class _T(GatewayEvent): + @classmethod + def create(cls, obj): + obj[field] = obj + self = cls(obj) + self.validate() + return self + + return _T class Ready(GatewayEvent): @@ -20,3 +37,131 @@ class Ready(GatewayEvent): session_id = skema.StringType() user = skema.ModelType(User) guilds = skema.ListType(skema.ModelType(Guild)) + + +class Resumed(GatewayEvent): + pass + + +class GuildCreate(Sub('guild')): + guild = skema.ModelType(Guild) + unavailable = skema.BooleanType(default=None) + + +class GuildUpdate(Sub('guild')): + guild = skema.ModelType(Guild) + + +class GuildDelete(GatewayEvent): + id = skema.SnowflakeType() + unavailable = skema.BooleanType(default=None) + + +class ChannelCreate(Sub('channel')): + channel = skema.ModelType(Channel) + + +class ChannelUpdate(Sub('channel')): + channel = skema.ModelType(Channel) + + +class ChannelDelete(Sub('channel')): + channel = skema.ModelType(Channel) + + +class ChannelPinsUpdate(GatewayEvent): + channel_id = skema.SnowflakeType() + last_pin_timestamp = skema.IntType() + + +class GuildBanAdd(Sub('user')): + user = skema.ModelType(User) + + +class GuildBanRemove(Sub('user')): + user = skema.ModelType(User) + + +class GuildEmojisUpdate(GatewayEvent): + pass + + +class GuildIntegrationsUpdate(GatewayEvent): + pass + + +class GuildMembersChunk(GatewayEvent): + guild_id = skema.SnowflakeType() + members = skema.ListType(skema.ModelType(GuildMember)) + + +class GuildMemberAdd(Sub('member')): + member = skema.ModelType(GuildMember) + + +class GuildMemberRemove(GatewayEvent): + guild_id = skema.SnowflakeType() + user = skema.ModelType(User) + + +class GuildMemberUpdate(GatewayEvent): + guild_id = skema.SnowflakeType() + user = skema.ModelType(User) + roles = skema.ListType(skema.SnowflakeType()) + + +class GuildRoleCreate(GatewayEvent): + guild_id = skema.SnowflakeType() + role = skema.ModelType(Role) + + +class GuildRoleUpdate(GatewayEvent): + guild_id = skema.SnowflakeType() + role = skema.ModelType(Role) + + +class GuildRoleDelete(GatewayEvent): + guild_id = skema.SnowflakeType() + role = skema.ModelType(Role) + + +class MessageCreate(Sub('message')): + message = skema.ModelType(Message) + + +class MessageUpdate(Sub('message')): + message = skema.ModelType(Message) + + +class MessageDelete(GatewayEvent): + id = skema.SnowflakeType() + channel_id = skema.SnowflakeType() + + +class MessageDeleteBulk(GatewayEvent): + channel_id = skema.SnowflakeType() + ids = skema.ListType(skema.SnowflakeType()) + + +class PresenceUpdate(GatewayEvent): + user = skema.ModelType(User) + guild_id = skema.SnowflakeType() + roles = skema.ListType(skema.SnowflakeType()) + game = skema.StringType() + status = skema.StringType() + + +class TypingStart(GatewayEvent): + channel_id = skema.SnowflakeType() + user_id = skema.SnowflakeType() + timestamp = skema.IntType() + + +class VoiceStateUpdate(Sub('state')): + state = skema.ModelType(VoiceState) + + +class VoiceServerUpdate(GatewayEvent): + token = skema.StringType() + endpoint = skema.StringType() + guild_id = skema.SnowflakeType() diff --git a/disco/state.py b/disco/state.py new file mode 100644 index 0000000..e69de29 diff --git a/disco/types/__init__.py b/disco/types/__init__.py index e69de29..5e6f73b 100644 --- a/disco/types/__init__.py +++ b/disco/types/__init__.py @@ -0,0 +1,5 @@ +from disco.types.channel import Channel +from disco.types.guild import Guild, GuildMember, Role +from disco.types.user import User +from disco.types.message import Message +from disco.types.voice import VoiceState diff --git a/disco/types/channel.py b/disco/types/channel.py new file mode 100644 index 0000000..4ae1465 --- /dev/null +++ b/disco/types/channel.py @@ -0,0 +1,42 @@ +import skema + +from holster.enum import Enum + +# from disco.types.guild import Guild +from disco.types.user import User + + +ChannelType = Enum( + GUILD_TEXT=0, + DM=1, + GUILD_VOICE=2, + GROUP_DM=3, +) + +PermissionOverwriteType = Enum( + ROLE='role', + MEMBER='member' +) + + +class PermissionOverwrite(skema.Model): + id = skema.SnowflakeType() + type = skema.StringType(choices=PermissionOverwriteType.ALL_VALUES) + + allow = skema.IntType() + deny = skema.IntType() + + +class Channel(skema.Model): + id = skema.SnowflakeType() + + name = skema.StringType() + topic = skema.StringType() + last_message_id = skema.SnowflakeType() + position = skema.IntType() + bitrate = skema.IntType(required=False) + + recipient = skema.ModelType(User, required=False) + type = skema.IntType(choices=ChannelType.ALL_VALUES) + + permission_overwrites = skema.ListType(skema.ModelType(PermissionOverwrite)) diff --git a/disco/types/guild.py b/disco/types/guild.py index 397ea57..8ff6835 100644 --- a/disco/types/guild.py +++ b/disco/types/guild.py @@ -1,7 +1,58 @@ import skema -from disco.util.oop import TypedClass +from disco.util.types import PreHookType +from disco.types.user import User +from disco.types.voice import VoiceState +from disco.types.channel import Channel + + +class Emoji(skema.Model): + id = skema.SnowflakeType() + name = skema.StringType() + require_colons = skema.BooleanType() + managed = skema.BooleanType() + roles = skema.ListType(skema.SnowflakeType()) + + +class Role(skema.Model): + id = skema.SnowflakeType() + name = skema.StringType() + hoist = skema.BooleanType() + managed = skema.BooleanType() + color = skema.IntType() + permissions = skema.IntType() + position = skema.IntType() + + +class GuildMember(skema.Model): + user = skema.ModelType(User) + mute = skema.BooleanType() + deaf = skema.BooleanType() + joined_at = PreHookType(lambda k: k[:-6], skema.DateTimeType()) + roles = skema.ListType(skema.SnowflakeType()) class Guild(skema.Model): id = skema.SnowflakeType() + + owner_id = skema.SnowflakeType() + afk_channel_id = skema.SnowflakeType() + embed_channel_id = skema.SnowflakeType() + + name = skema.StringType() + icon = skema.BinaryType(None) + splash = skema.BinaryType(None) + region = skema.StringType() + + afk_timeout = skema.IntType() + embed_enabled = skema.BooleanType() + verification_level = skema.IntType() + mfa_level = skema.IntType() + + features = skema.ListType(skema.StringType()) + + members = skema.ListType(skema.ModelType(GuildMember)) + voice_states = skema.ListType(skema.ModelType(VoiceState)) + channels = skema.ListType(skema.ModelType(Channel)) + roles = skema.ListType(skema.ModelType(Role)) + emojis = skema.ListType(skema.ModelType(Emoji)) diff --git a/disco/types/message.py b/disco/types/message.py new file mode 100644 index 0000000..534f299 --- /dev/null +++ b/disco/types/message.py @@ -0,0 +1,44 @@ +import skema + +from disco.util.types import PreHookType +from disco.types.user import User + + +class MessageEmbed(skema.Model): + title = skema.StringType() + type = skema.StringType() + description = skema.StringType() + url = skema.StringType() + + +class MessageAttachment(skema.Model): + id = skema.SnowflakeType() + filename = skema.StringType() + url = skema.StringType() + proxy_url = skema.StringType() + size = skema.IntType() + height = skema.IntType() + width = skema.IntType() + + +class Message(skema.Model): + id = skema.SnowflakeType() + channel_id = skema.SnowflakeType() + + author = skema.ModelType(User) + content = skema.StringType() + nonce = skema.StringType() + + timestamp = PreHookType(lambda k: k[:-6], skema.DateTimeType()) + edited_timestamp = PreHookType(lambda k: k[:-6], skema.DateTimeType()) + + tts = skema.BooleanType() + mention_everyone = skema.BooleanType() + + pinned = skema.BooleanType(required=False) + + mentions = skema.ListType(skema.ModelType(User)) + mention_roles = skema.ListType(skema.SnowflakeType()) + + embeds = skema.ListType(skema.ModelType(MessageEmbed)) + attachment = skema.ListType(skema.ModelType(MessageAttachment)) diff --git a/disco/types/user.py b/disco/types/user.py index 5411d15..f1936f5 100644 --- a/disco/types/user.py +++ b/disco/types/user.py @@ -3,3 +3,10 @@ import skema class User(skema.Model): id = skema.SnowflakeType() + + username = skema.StringType() + discriminator = skema.StringType() + avatar = skema.BinaryType(None) + + verified = skema.BooleanType(required=False) + email = skema.EmailType(required=False) diff --git a/disco/types/voice.py b/disco/types/voice.py new file mode 100644 index 0000000..91b1ea5 --- /dev/null +++ b/disco/types/voice.py @@ -0,0 +1,5 @@ +import skema + + +class VoiceState(skema.Model): + id = skema.SnowflakeType() diff --git a/disco/util/types.py b/disco/util/types.py new file mode 100644 index 0000000..d8b6d0d --- /dev/null +++ b/disco/util/types.py @@ -0,0 +1,18 @@ +from skema import BaseType + + +class PreHookType(BaseType): + _hashable = False + + def __init__(self, func, field, **kwargs): + self.func = func + self.field = field + + super(PreHookType, self).__init__(**kwargs) + + def to_python(self, value): + value = self.func(value) + return self.field.to_python(value) + + def to_storage(self, *args, **kwargs): + return self.field.to_storage(*args, **kwargs)