Exemplo n.º 1
0
class GetContacts:
    def __init__(self, print_fn=print, loop=None):
        self._print_fn = print_fn
        self._driver = Whalesong(profile=path.join(path.dirname(__file__),
                                                   'profile'),
                                 loadstyles=True,
                                 loop=loop)

    @property
    def loop(self):
        return self._driver.loop

    def echo(self, txt):
        self._print_fn(txt)

    async def check_stream(self):
        stream = await self._driver.stream.get_model()
        self.echo("Stream: {}".format(stream.stream))
        self.echo("State: {}".format(stream.state))

    async def monitor_stream(self):
        self.echo('Monitor stream')
        contact_it = None
        async for evt in self._driver.stream.monitor_field('stream'):
            self.echo('Stream value: {}'.format(evt['value']))

            if evt['value'] == Stream.Stream.CONNECTED:
                if contact_it is None:
                    contact_it = self._driver.contacts.get_items()
                    ensure_future(self.list_contacts(contact_it))
            else:
                if contact_it is not None:
                    await self._driver.cancel_iterators()
                    contact_it = None

    async def list_contacts(self, it):
        self.echo('List contacts')
        async for contact in it:
            self.echo('Contact: {}'.format(contact))

        self.echo('List contacts finished')

    async def start(self):
        await self._driver.start()

        loop = get_event_loop()
        loop.add_signal_handler(
            SIGINT, lambda *args: ensure_future(self._driver.stop()))
        loop.add_signal_handler(
            SIGTERM, lambda *args: ensure_future(self._driver.stop()))

        ensure_future(self.check_stream()),
        ensure_future(self.monitor_stream())

        await self._driver.wait_until_stop()
class PresenceMonitor:

    def __init__(self, print_fn=print, loop=None):
        self._print_fn = print_fn

        self._driver = Whalesong(
            driver=WhalesongDriver(profile=path.join(path.dirname(__file__), 'profile-chromium'),
                                   headless=True,
                                   loop=loop),
            loop=loop
        )

    @property
    def loop(self):
        return self._driver.loop

    def echo(self, txt):
        self._print_fn(txt)

    async def check_stream(self):
        stream = await self._driver.stream.get_model()
        self.echo("Stream: {}".format(stream.stream))
        self.echo("State: {}".format(stream.state))

    async def monitor_stream(self):
        self.echo('Monitor stream')
        chat_it = None
        async for evt in self._driver.stream.monitor_field('stream'):
            self.echo('Stream value: {}'.format(evt['value']))

            if evt['value'] == Stream.Stream.CONNECTED:
                self._driver.display_info.set_available_permanent()

                if chat_it is None:
                    chat_it = self._driver.chats.get_items()
                    ensure_future(self.start_monitor_presences(chat_it))
            else:
                if chat_it is not None:
                    await self._driver.cancel_iterators()
                    chat_it = None

    async def start_monitor_presences(self, it):
        self.echo('Listing chats...')
        async for chat in it:
            self.echo('Starting: {}'.format(chat.name))
            ensure_future(self.monitor_chat(chat))

        self.echo('List chats finished')

    async def monitor_chat(self, chat):
        if chat.is_group:
            self.echo(f'Chat `{chat.name}` is group')
            await self._driver.chats[chat.id].presence.subscribe()

            async for chat_state in self._driver.chats[chat.id].presence.chat_states.get_items():
                ensure_future(self.monitor_chat_state(
                    chat,
                    self._driver.chats[chat.id].presence.chat_states[chat_state.id])
                )

            async for chat_state in self._driver.chats[chat.id].presence.chat_states.monitor_add():
                ensure_future(self.monitor_chat_state(
                    chat,
                    self._driver.chats[chat.id].presence.chat_states[chat_state.id])
                )

        else:
            ensure_future(self.monitor_chat_state(
                chat,
                self._driver.chats[chat.id].presence.chat_state
            ))

    async def monitor_chat_state(self, chat, chat_state_manager):
        chat_state = await chat_state_manager.get_model()

        def echo_state(type):
            self.echo(f'User {chat_state.id} is {type} in chat {chat.name} ({chat.id})')

        echo_state(chat_state.type)

        async for evt in chat_state_manager.monitor_field('type'):
            echo_state(evt['value'])

    async def start(self):
        await self._driver.start()

        loop = get_event_loop()
        loop.add_signal_handler(SIGINT, lambda *args: ensure_future(self._driver.stop()))
        loop.add_signal_handler(SIGTERM, lambda *args: ensure_future(self._driver.stop()))

        ensure_future(self.check_stream()),
        ensure_future(self.monitor_stream())

        await self._driver.wait_until_stop()
Exemplo n.º 3
0
class GetMessages:
    def __init__(self, print_fn=print, loop=None):
        self._print_fn = print_fn
        self._driver = Whalesong(profile=path.join(path.dirname(__file__),
                                                   'profile'),
                                 loadstyles=True,
                                 loop=loop)

        try:
            mkdir(OUTPUT_DIR)
        except FileExistsError:
            pass

    @property
    def loop(self):
        return self._driver.loop

    def echo(self, txt):
        self._print_fn(txt)

    async def check_stream(self):
        stream = await self._driver.stream.get_model()
        self.echo("Stream: {}".format(stream.stream))
        self.echo("State: {}".format(stream.state))

    async def monitor_stream(self):
        self.echo('Monitor stream')
        messages_it = None
        new_message_monitor = None
        message_ack_monitor = None

        async for evt in self._driver.stream.monitor_field('stream'):
            self.echo('Stream value: {}'.format(evt['value']))

            if evt['value'] == Stream.Stream.CONNECTED:
                if messages_it is None:
                    messages_it = self._driver.messages.get_items()
                    ensure_future(self.list_messages(messages_it))

                if new_message_monitor is None:
                    new_message_monitor = self._driver.messages.monitor_new()
                    ensure_future(
                        self.monitor_new_messages(new_message_monitor))

                if message_ack_monitor is None:
                    message_ack_monitor = self._driver.messages.monitor_field(
                        'ack')
                    ensure_future(
                        self.monitor_message_acks(message_ack_monitor))
            else:
                if messages_it is not None:
                    messages_it = None
                    await self._driver.cancel_iterators()

                if new_message_monitor is not None:
                    self._driver.stop_monitor(new_message_monitor)
                    new_message_monitor = None

                if message_ack_monitor is not None:
                    self._driver.stop_monitor(message_ack_monitor)
                    message_ack_monitor = None

    async def list_messages(self, it):
        self.echo('List messages')
        async for message in it:
            self.echo('Message: {}'.format(message))

            if isinstance(message, MediaMixin):
                await self.download_media(message)
            elif isinstance(message, LocationMessage):
                await self.store_thumbnail(message.id, message.body)
            elif isinstance(message, TextMessage) and message.thumbnail:
                await self.store_thumbnail(message.id, message.thumbnail)

        self.echo('List messages finished')

    async def monitor_new_messages(self, it):
        self.echo('Monitor new messages')
        async for message in it:
            self.echo('New message: {}'.format(message))

            if isinstance(message, MediaMixin):
                await self.download_media(message)
            elif isinstance(message, LocationMessage):
                await self.store_thumbnail(message.id, message.body)
            elif isinstance(message, TextMessage) and message.thumbnail:
                await self.store_thumbnail(message.id, message.thumbnail)

        self.echo('Stop new messages bot')

    async def monitor_message_acks(self, it):
        self.echo('Monitor message ack')
        async for ack in it:
            self.echo('ACK: {}'.format(ack))

        self.echo('Stop message acks bot')

    async def start(self):
        await self._driver.start()

        loop = get_event_loop()
        loop.add_signal_handler(
            SIGINT, lambda *args: ensure_future(self._driver.stop()))
        loop.add_signal_handler(
            SIGTERM, lambda *args: ensure_future(self._driver.stop()))

        ensure_future(self.check_stream()),
        ensure_future(self.monitor_stream())

        await self._driver.wait_until_stop()

    async def download_media(self, message):
        if isinstance(message, MediaFrameMixin) and not isinstance(
                message, StickerMessage):
            await self.store_thumbnail(message.id, message.body)

        media_data = await self._driver.messages.download_media(message)

        mimetype = message.mimetype

        if ';' in mimetype:
            mimetype, _ = mimetype.split(';', 1)

        ext = mimetypes.guess_extension(mimetype, strict=False)

        await self._store_media('{}{}'.format(message.id, ext or '.bin'),
                                media_data.read())

    async def store_thumbnail(self, message_id, image_data):
        await self._store_media('{}_thumb.jpg'.format(message_id), image_data)

    async def _store_media(self, filename, filedata):
        filepath = path.join(OUTPUT_DIR, filename)

        self.echo('Storing file {}'.format(filepath))

        with open(filepath, 'wb') as f:
            f.write(filedata)
Exemplo n.º 4
0
class StatusMonitor:

    def __init__(self, print_fn=print, loop=None):
        self._print_fn = print_fn
        self._driver = Whalesong(
            driver=WhalesongDriver(profile=path.join(path.dirname(__file__), 'profile-chromium'),
                                   headless=True,
                                   loop=loop),
            loop=loop
        )

    @property
    def loop(self):
        return self._driver.loop

    def echo(self, txt):
        self._print_fn(txt)

    async def check_stream(self):
        stream = await self._driver.stream.get_model()
        self.echo("Stream: {}".format(stream.stream))
        self.echo("State: {}".format(stream.state))

    async def check_conn(self):
        stream = await self._driver.conn.get_model()
        self.echo("REF: {}".format(stream.ref))
        self.echo("Battery: {}".format(stream.battery))

    async def check_storage(self):
        storage = await self._driver.storage.get_storage()
        self.echo("Storage: {}".format(storage))

    async def monitor_stream(self):
        self.echo('Monitor stream')
        async for evt in self._driver.stream.monitor_field('stream'):
            self.echo('Stream value: {}'.format(evt['value']))
            img = await self._driver.screenshot()
            with open(path.join(OUTPUT_DIR, 'screenshot.png'), 'wb') as f:
                f.write(img.read())

            self.echo('Screenshot saved!')

    async def monitor_state(self):
        self.echo('Monitor state')
        async for evt in self._driver.stream.monitor_field('state'):
            self.echo(evt)
            self.echo('State value: {}'.format(evt['value']))
            if evt['value'] == Stream.State.UNPAIRED_IDLE:
                self.echo('Refreshing QR')
                await self._driver.stream.poke()
            elif evt['value'] == Stream.State.CONFLICT:
                self.echo('Taking over...')
                await self._driver.stream.takeover()

    async def monitor_ref(self):
        self.echo('Monitor ref')
        async for evt in self._driver.conn.monitor_field('ref'):
            self.echo('New REF value: {}'.format(evt['value']))

            try:
                img = await self._driver.qr()
                with open(path.join(OUTPUT_DIR, 'qr.png'), 'wb') as f:
                    f.write(img.read())

                self.echo('QR saved!')
            except Exception as ex:
                self.echo(ex)
                self.echo('Error getting qr')

    async def monitor_battery(self):
        self.echo('Monitor battery')
        async for evt in self._driver.conn.monitor_field('battery'):
            self.echo('Battery level: {}'.format(evt['value']))

    async def monitor_storage(self):
        self.echo('Monitor storage')
        async for evt in self._driver.storage.monitor_item_storage():
            self.echo('New storage item: {}'.format(evt))

    async def start(self):
        await self._driver.start()

        loop = get_event_loop()
        loop.add_signal_handler(SIGINT, lambda *args: ensure_future(self._driver.stop()))
        loop.add_signal_handler(SIGTERM, lambda *args: ensure_future(self._driver.stop()))

        ensure_future(self.check_stream())
        ensure_future(self.check_conn())
        ensure_future(self.check_storage())
        ensure_future(self.monitor_stream())
        ensure_future(self.monitor_state())
        ensure_future(self.monitor_ref())
        ensure_future(self.monitor_battery())
        ensure_future(self.monitor_storage())

        await self._driver.wait_until_stop()
Exemplo n.º 5
0
class GetStatuses:
    def __init__(self, print_fn=print, loop=None):
        self._print_fn = print_fn
        self._driver = Whalesong(driver=WhalesongDriver(profile=path.join(
            path.dirname(__file__), 'profile-chromium'),
                                                        headless=False,
                                                        loop=loop),
                                 loop=loop)

        try:
            mkdir(OUTPUT_DIR)
        except FileExistsError:
            pass

    @property
    def loop(self):
        return self._driver.loop

    def echo(self, text):
        self._print_fn(text)

    async def check_stream(self):
        stream = await self._driver.stream.get_model()
        self.echo("Stream: {}".format(stream.stream))
        self.echo("State: {}".format(stream.state))

    async def monitor_stream(self):
        self.echo('Monitor stream')
        statuses_it = None

        async for evt in self._driver.stream.monitor_field('stream'):
            self.echo('Stream value: {}'.format(evt['value']))

            if evt['value'] == Stream.Stream.CONNECTED:
                if statuses_it is None:
                    await self._driver.status_v3.sync()
                    statuses_it = self._driver.status_v3.get_unexpired(
                        unread=False)
                    ensure_future(self.list_unread_statuses(statuses_it))

            else:
                if statuses_it is not None:
                    statuses_it = None
                    await self._driver.cancel_iterators()

    async def list_unread_statuses(self, it):
        self.echo('List statuses')

        async for status in it:
            print('Seeing all statuses from [{}]'.format(status.id))
            async for msg in self._driver.status_v3[status.id].get_submanager(
                    'msgs').get_items():
                await self._driver.status_v3[status.id].send_read_status(msg.id
                                                                         )
                if isinstance(msg, MediaMixin):
                    self.echo('Saving status media from [{}]'.format(
                        status.id))
                    await self.download_media(msg)

    async def start(self):
        await self._driver.start()

        loop = get_event_loop()
        loop.add_signal_handler(
            SIGINT, lambda *args: ensure_future(self._driver.stop()))
        loop.add_signal_handler(
            SIGTERM, lambda *args: ensure_future(self._driver.stop()))

        ensure_future(self.check_stream())
        ensure_future(self.monitor_stream())

        await self._driver.wait_until_stop()

    async def download_media(self, message):
        if isinstance(message, MediaFrameMixin) and not isinstance(
                message, StickerMessage):
            await self.store_thumbnail(message.id, message.body)

        media_data = await self._driver.messages.download_media(message)

        mimetype = message.mimetype

        if ';' in mimetype:
            mimetype, _ = mimetype.split(';', 1)

        ext = mimetypes.guess_extension(mimetype, strict=False)

        await self._store_media('{}{}'.format(message.id, ext or '.bin'),
                                media_data.read())

    async def store_thumbnail(self, message_id, image_data):
        await self._store_media('{}_thumb.jpg'.format(message_id), image_data)

    async def _store_media(self, filename, filedata):
        filepath = path.join(OUTPUT_DIR, filename)

        self.echo('Storing file {}'.format(filepath))

        with open(filepath, 'wb') as f:
            f.write(filedata)
Exemplo n.º 6
0
class Minibot:
    def __init__(self, print_fn=print, loop=None):
        self._print_fn = print_fn
        self._driver = Whalesong(profile=path.join(path.dirname(__file__),
                                                   'profile'),
                                 loadstyles=True,
                                 loop=loop)

    @property
    def loop(self):
        return self._driver.loop

    def echo(self, txt):
        self._print_fn(txt)

    async def download_url(self, url):
        async with ClientSession() as session:
            async with session.get(url) as resp:
                content_type = resp.headers.get(hdrs.CONTENT_TYPE)
                data = BytesIO(await resp.read())

        if content_type and ';' in content_type:
            content_type = content_type[:content_type.index(';')]

        return content_type, data

    async def monitor_stream(self):
        self.echo('Monitor stream')
        new_message_monitor = None

        async for evt in self._driver.stream.monitor_field('stream'):
            self.echo('Stream value: {}'.format(evt['value']))

            if evt['value'] == Stream.Stream.CONNECTED:
                if new_message_monitor is None:
                    new_message_monitor = self._driver.messages.monitor_new()
                    ensure_future(
                        self.monitor_new_messages(new_message_monitor))
            else:
                if new_message_monitor is not None:
                    self._driver.stop_monitor(new_message_monitor)
                    new_message_monitor = None

    async def monitor_new_messages(self, it):
        self.echo('Monitor new messages')
        async for message in it:
            try:
                if message.type == MessageTypes.CHAT:
                    if message.body.startswith('/echo '):
                        ensure_future(self.make_echo(message))
                    elif message.body.startswith('/contact '):
                        ensure_future(self.make_contact(message))
                    elif message.body.startswith('/download '):
                        ensure_future(self.make_download(message))
                    elif message.body.startswith('/send '):
                        ensure_future(self.make_message(message))
                    elif message.body.startswith('/link '):
                        ensure_future(self.make_link(message))
                    elif message.body.startswith('/exist '):
                        ensure_future(self.make_exists(message))
                    elif message.body.startswith('/sticker '):
                        ensure_future(self.make_sticker(message))
                    elif message.body.startswith('/status '):
                        ensure_future(self.make_status(message))
                    elif message.body.startswith('/pushname '):
                        ensure_future(self.make_pushname(message))

            except Exception as ex:
                self.echo('Ignoring message {} because error : {}'.format(
                    message.id, ex))

        self.echo('Stop new messages bot')

    async def make_echo(self, message):
        self._driver.chats[message.chat.id].send_seen()
        self.echo(message.chat.id)
        text = message.body[len('/echo '):].strip()
        self.echo('Sent message: {}'.format(await self._driver.chats[
            message.chat.id].send_text(text, message.id)))

    async def make_contact(self, message):
        self._driver.chats[message.chat.id].send_seen()
        contact_id = message.body[len('/contact '):].strip()

        contact = await self._driver.contacts.get_item_by_id(contact_id)

        self.echo('Sent message (by phone): {}'.format(
            await self._driver.chats[message.chat.id].send_contact_phone(
                contact.formatted_name, contact.userid, message.id)))

        self.echo('Sent message (by contact id): {}'.format(
            await self._driver.chats[message.chat.id
                                     ].send_contact(contact_id, message.id)))

    async def make_download(self, message):
        self._driver.chats[message.chat.id].send_seen()
        url = message.body[len('/download '):].strip()

        content_type, data = await self.download_url(url)

        filename = url
        if '?' in filename:
            filename = filename[:filename.index('?')]
        if '#' in filename:
            filename = filename[:filename.index('#')]
        _, filename = filename.rstrip('/').rsplit('/', 1)

        self.echo('Sent message: {}'.format(await self._driver.chats[
            message.chat.id].send_media(data, content_type, filename, url,
                                        message.id)))

    async def make_message(self, message):
        self._driver.chats[message.chat.id].send_seen()
        self.echo(message.chat.id)
        cmd = message.body[len('/send '):].strip()
        contact_id, text = cmd.split(' ')
        if not contact_id.endswith('@c.us'):
            contact_id = f'{contact_id.strip()}@c.us'

        chat = await self._driver.chats.ensure_chat_with_contact(
            contact_id=contact_id)
        msg = await self._driver.chats[chat.id].send_text(text)
        self.echo('Sent message: {}'.format(msg))

        await self._driver.chats[message.chat.id
                                 ].send_text(f'Message sent to {chat.id}')

    async def make_link(self, message):
        self._driver.chats[message.chat.id].send_seen()
        self.echo(message.chat.id)
        text = message.body[len('/link '):].strip()
        self.echo('Sent message: {}'.format(await self._driver.chats[
            message.chat.id].send_text(text)))

    async def make_exists(self, message):
        self._driver.chats[message.chat.id].send_seen()

        contact_id = message.body[len('/exist '):].strip()

        if not contact_id.endswith('@c.us'):
            contact_id = f'{contact_id.strip()}@c.us'

        exists = await self._driver.wap.query_exist(contact_id=contact_id)
        self.echo('Sent message: {}'.format(await self._driver.chats[
            message.chat.id
        ].send_text('It exists' if exists else 'It does not exist',
                    message.id)))

    async def make_sticker(self, message):
        self._driver.chats[message.chat.id].send_seen()

        sticker_pack_name = message.body[len('/sticker '):].strip()

        if sticker_pack_name == 'list':
            ensure_future(self.make_sticker_pack_list(message))
            return

        await self._driver.sticker_packs.fetch_all_pages()

        try:
            sticker_pack = await self._driver.sticker_packs.get_item_by_name(
                sticker_pack_name)
        except ModuleNotFoundError:
            self.echo('Sent message: {}'.format(await self._driver.chats[
                message.chat.id
            ].send_text(f'Sticker pack "{sticker_pack_name}" does not exist',
                        message.id)))

            return

        await self._driver.sticker_packs[sticker_pack.id].stickers.fetch()

        sticker = choice([
            s async for s in self._driver.sticker_packs[
                sticker_pack.id].stickers.get_items()
        ])

        msg = await self._driver.sticker_packs[sticker_pack.id].stickers[
            sticker.id].send_to_chat(chat_id=message.chat.id,
                                     quoted_msg_id=message.id)

        self.echo(f'Sent message: {msg}')

    async def make_sticker_pack_list(self, message):
        await self._driver.sticker_packs.fetch_all_pages()

        async for sticker_pack in self._driver.sticker_packs.get_items():
            content_type, image = await self.download_url(sticker_pack.url)

            self.echo('Sent message: {}'.format(await self._driver.chats[
                message.chat.id].send_media(image,
                                            content_type,
                                            caption=sticker_pack.name)))

    async def make_status(self, message):
        new_status = message.body[len('/status '):].strip()
        await self._driver.status.set_my_status(new_status=new_status)
        self.echo(f'Set status: {new_status}')

    async def make_pushname(self, message):
        new_pushname = message.body[len('/pushname '):].strip()
        await self._driver.conn.update_pushname(name=new_pushname)
        self.echo(f'Set pushname: {new_pushname}')

    async def start(self):
        await self._driver.start()

        loop = get_event_loop()
        loop.add_signal_handler(
            SIGINT, lambda *args: ensure_future(self._driver.stop()))
        loop.add_signal_handler(
            SIGTERM, lambda *args: ensure_future(self._driver.stop()))

        ensure_future(self.monitor_stream())

        await self._driver.wait_until_stop()
class GetLiveLocations:
    def __init__(self, print_fn=print, loop=None):
        self._print_fn = print_fn
        self._driver = Whalesong(driver=WhalesongDriver(profile=path.join(
            path.dirname(__file__), 'profile-chromium'),
                                                        headless=True,
                                                        loop=loop),
                                 loop=loop)

    @property
    def loop(self):
        return self._driver.loop

    def echo(self, txt):
        self._print_fn(txt)

    async def check_stream(self):
        stream = await self._driver.stream.get_model()
        self.echo("Stream: {}".format(stream.stream))
        self.echo("State: {}".format(stream.state))

    async def monitor_stream(self):
        self.echo('Monitor stream')
        live_loc_it = None
        new_message_monitor = None
        message_ack_monitor = None

        async for evt in self._driver.stream.monitor_field('stream'):
            self.echo('Stream value: {}'.format(evt['value']))

            if evt['value'] == Stream.Stream.CONNECTED:

                if live_loc_it is None:
                    live_loc_it = self._driver.live_locations.get_items()
                    ensure_future(self.list_live_locations(live_loc_it))

                if new_message_monitor is None:
                    new_live_locations = self._driver.live_locations.monitor_add(
                    )
                    ensure_future(
                        self.monitor_new_live_locations(new_live_locations))

            else:
                if live_loc_it is not None:
                    live_loc_it = None
                    await self._driver.cancel_iterators()

                if new_message_monitor is not None:
                    self._driver.stop_monitor(new_message_monitor)
                    new_message_monitor = None

                if message_ack_monitor is not None:
                    self._driver.stop_monitor(message_ack_monitor)
                    message_ack_monitor = None

    async def list_live_locations(self, it):
        self.echo('List live locations')
        async for live_location in it:
            self.echo('Live location: {}'.format(live_location))
            ensure_future(self.monitor_live_location(live_location))

        self.echo('List messages finished')

    async def monitor_new_live_locations(self, it):
        self.echo('Monitor live locations')
        async for live_location in it:
            self.echo('New live location: {}'.format(live_location))
            ensure_future(self.monitor_live_location(live_location))

    async def monitor_live_location(self, live_location):
        self._driver.live_locations[live_location.id].subscribe()

        for participant in live_location.participants:
            ensure_future(self.monitor_participant(live_location, participant))

        async for participant in self._driver.live_locations[
                live_location.id].participants.monitor_add():
            ensure_future(self.monitor_participant(live_location, participant))

    async def monitor_participant(self, live_location, participant):
        async for new_part in self._driver.live_locations[live_location.id] \
                .participants[participant.id].monitor_model():
            self.echo(
                'Participant at live location `{}` now are in {}@{}'.format(
                    live_location.id, new_part.lat, new_part.lng))

    async def start(self):
        await self._driver.start()

        loop = get_event_loop()
        loop.add_signal_handler(
            SIGINT, lambda *args: ensure_future(self._driver.stop()))
        loop.add_signal_handler(
            SIGTERM, lambda *args: ensure_future(self._driver.stop()))

        ensure_future(self.check_stream()),
        ensure_future(self.monitor_stream())

        await self._driver.wait_until_stop()
Exemplo n.º 8
0
class GetStickers:

    def __init__(self, print_fn=print, loop=None):
        self._print_fn = print_fn
        self._driver = Whalesong(
            profile=path.join(path.dirname(__file__), 'profile'),
            loadstyles=True,
            loop=loop
        )

        try:
            mkdir(OUTPUT_DIR)
        except FileExistsError:
            pass

    @property
    def loop(self):
        return self._driver.loop

    def echo(self, txt):
        self._print_fn(txt)

    async def check_stream(self):
        stream = await self._driver.stream.get_model()
        try:
            self.echo("Stream: {}".format(stream.stream))
            self.echo("State: {}".format(stream.state))
        except AttributeError:
            self.echo("Error: {}".format(stream))

    async def monitor_stream(self):
        self.echo('Monitor stream')

        async for evt in self._driver.stream.monitor_field('stream'):
            self.echo('Stream value: {}'.format(evt['value']))

            if evt['value'] == Stream.Stream.CONNECTED:
                ensure_future(self.fetch_sticker_packs())

    async def fetch_sticker_packs(self):
        self.echo('Fetching all sticker packs')
        await self._driver.sticker_packs.fetch_all_pages()

        async for sticker_pack in self._driver.sticker_packs.get_items():
            await self._driver.sticker_packs[sticker_pack.id].stickers.fetch()
            await self.download_all_sticker(sticker_pack)

    async def download_all_sticker(self, sticker_pack):
        self.echo(f'Download all stickers from sticker pack {sticker_pack.name}')
        async for sticker in self._driver.sticker_packs[sticker_pack.id].stickers.get_items():
            self.echo(f'Download sticker {sticker.id} from sticker pack {sticker_pack.name}')
            await self.download_media(sticker_pack, sticker)

        self.echo(f'All stickers download from sticker pack {sticker_pack.name}')

    async def start(self):
        await self._driver.start()

        loop = get_event_loop()
        loop.add_signal_handler(SIGINT, lambda *args: ensure_future(self._driver.stop()))
        loop.add_signal_handler(SIGTERM, lambda *args: ensure_future(self._driver.stop()))

        ensure_future(self.check_stream()),
        ensure_future(self.monitor_stream())

        await self._driver.wait_until_stop()

    async def download_media(self, sticker_pack, sticker):
        media_data = await self._driver.sticker_packs[sticker_pack.id].stickers[sticker.id].download_image()

        ext = mimetypes.guess_extension(sticker.mimetype, strict=False)

        sticker_id = sticker.id.replace('/', '_').replace('\\', '_').replace('=', '_')

        await self._store_media(f'{sticker_pack.name}-{sticker_id}{ext}', media_data.read())

    async def _store_media(self, filename, filedata):
        filepath = path.join(OUTPUT_DIR, filename)

        self.echo('Storing file {}'.format(filepath))

        with open(filepath, 'wb') as f:
            f.write(filedata)