initial commit

kurocore/utils/database/database.py

@@ -0,0 +1,29 @@
+import logging
+from asyncio import get_event_loop
+from peewee_async import PostgresqlDatabase, Manager
+
+from kurocore.logger import BotLogger
+
+
+class Database:
+    db: PostgresqlDatabase = PostgresqlDatabase(None)
+    manager: Manager
+
+    def __init__(self, settings):
+        db_settings = settings.database
+
+        driver_name = db_settings.driver
+        host = db_settings.host
+        name = db_settings.name
+        user = db_settings.user
+        password = db_settings.password
+        port = db_settings.port
+
+        if driver_name.lower() not in ('postgres', 'postgresql', 'psql'):
+            raise
+
+        Database.db.init(name, user=user, password=password, host=host, port=port)
+        loop = get_event_loop()
+        Database.manager = Manager(Database.db, loop=loop)
+        Database.manager.allow_sync = logging.ERROR
+        BotLogger.log.info('Connected to database')

kurocore/utils/database/models.py

@@ -0,0 +1,173 @@
+from peewee import Model, IntegerField, BigIntegerField, DecimalField, TextField, ForeignKeyField, DateTimeField, \
+    BooleanField, CompositeKey, CharField, AutoField
+
+# from plugins import models
+from utils.database.database import Database
+
+
+class BaseModel(Model):
+    class Meta:
+        database = Database.db
+
+
+class Groups(BaseModel):
+    id = IntegerField(null=False)
+    name = TextField(default='')
+    is_vip = BooleanField(default=False)
+    is_admin = BooleanField(default=False)
+    is_tester = BooleanField(default=False)
+    multiplier = DecimalField(default=1)
+    sale = DecimalField(default=1)
+    max_multigen = IntegerField(default=3)
+
+
+class Works(BaseModel):
+    id = IntegerField(null=False)
+    name = TextField(default='')
+    required_lvl = IntegerField(default=1)
+    money_income = DecimalField(default=0)
+    min_exp = IntegerField(default=0)
+    max_exp = IntegerField(default=0)
+
+
+class Fractions(BaseModel):
+    id = IntegerField()
+    name = TextField(default='Без названия', null=False)
+    owner_id = IntegerField(null=False)
+    money = DecimalField()
+    exp = IntegerField()
+    level = IntegerField()
+
+
+class Auto(BaseModel):
+    id = IntegerField()
+    name = TextField()
+    price = DecimalField()
+
+    class Meta:
+        table_name = 'shop_auto'
+
+
+class Business(BaseModel):
+    id = IntegerField()
+    name = TextField()
+    price = DecimalField()
+    income = DecimalField()
+
+    class Meta:
+        table_name = 'shop_business'
+
+
+class Maid(BaseModel):
+    id = IntegerField()
+    name = TextField()
+    price = DecimalField()
+    income = DecimalField()
+
+    class Meta:
+        table_name = 'shop_maid'
+
+
+class Miner(BaseModel):
+    id = IntegerField()
+    name = TextField()
+    price = DecimalField()
+    income = DecimalField()
+
+    class Meta:
+        table_name = 'shop_miner'
+
+
+class Users(BaseModel):
+    id = IntegerField()
+    user_id = BigIntegerField()
+
+    name = TextField()
+    greeting = TextField()
+    pair = IntegerField()
+    donat = IntegerField()
+
+    balance = DecimalField(default=10000, max_digits=20)
+    invested = DecimalField(default=0, max_digits=20)
+    btc = DecimalField(default=0, max_digits=20, decimal_places=6)
+    level = IntegerField()
+    exp = DecimalField(default=0)
+
+    group_id = IntegerField()
+    group = ForeignKeyField(Groups)
+    fraction = ForeignKeyField(Fractions, null=True, lazy_load=True)
+
+    work = ForeignKeyField(Works)
+    work_time = DateTimeField()
+
+    income_time = DateTimeField()
+
+    auto = ForeignKeyField(Auto)
+    business = ForeignKeyField(Business)
+    maid = ForeignKeyField(Maid)
+    miner = ForeignKeyField(Miner)
+
+    subscribed = BooleanField()
+
+
+class Reports(BaseModel):
+    id = IntegerField(unique=True, null=False, primary_key=True)
+    text = TextField(null=False)
+    admin_answer = TextField(null=True)
+    from_id = BigIntegerField(null=False)
+    date = DateTimeField()
+    status = CharField(null=False)
+    attachments = TextField(null=True)
+
+
+class FractionMember(BaseModel):
+    fraction = ForeignKeyField(Fractions)
+    user = ForeignKeyField(Users)
+    is_moderator = BooleanField()
+    is_admin = BooleanField()
+
+    class Meta:
+        table_name = 'fraction_member'
+        primary_key = CompositeKey('fraction', 'user')
+
+
+# WD
+class WDModels(BaseModel):
+    model_id = IntegerField(null=False)
+    name = TextField(null=False)
+    image = TextField(null=False)
+    description = TextField(null=False)
+
+    class Meta:
+        table_name = 'wd_models'
+
+
+class WDUsers(BaseModel):
+    user_id = IntegerField(null=False)
+    model = ForeignKeyField(WDModels)
+    model_id = IntegerField(null=False, default=1)
+    orientation = TextField(null=False, default='PORTRAIT')
+
+    class Meta:
+        table_name = 'wd_users'
+        primary_key = CompositeKey('user_id', 'model_id')
+
+
+class WDLoras(BaseModel):
+    id = IntegerField(null=False)
+    name = TextField(null=False, default='NOT SET')
+    trigger_prompt = TextField(null=False, default='')
+    prompt_mixin = TextField(null=False, default='')
+
+    class Meta:
+        table_name = 'wd_loras'
+        primary_key = CompositeKey('id', 'trigger_prompt')
+
+
+class WDPrompts(BaseModel):
+    id = AutoField(null=False)
+    user_id = IntegerField(null=False)
+    prompt = TextField(null=False)
+
+    class Meta:
+        table_name = 'wd_prompts'

kurocore/utils/utils.py

@@ -0,0 +1,42 @@
+import re
+
+
+def chunk_array(array, chunk_number) -> list:
+    if len(array) <= chunk_number:
+        return [array]
+    out, tmp, index = [], [], 0
+    for el in array:
+        tmp.append(el)
+        index += 1
+        if index == chunk_number:
+            out.append(tmp)
+            tmp = []
+            index = 0
+    out.append(tmp)
+    return out
+
+
+def parse_attachments(attachments) -> tuple:
+    out = []
+    for attach in attachments:
+        t = attach['type']
+        if t == 'wall':
+            out.append(f'{t}{attach[t]["from_id"]}_{attach[t]["id"]}')
+        else:
+            out.append(f'{t}{attach[t]["owner_id"]}_{attach[t]["id"]}')
+    return tuple(out)
+
+
+def parse_mention(mention: (str, None)) -> int:
+    if not mention:
+        return 0
+    reg = r'\[id(\d+)\|.+\]'
+    match = re.match(reg, mention)
+    if not match:
+        return 0
+    return int(match.group(1))
+
+
+async def get_user_sex(user_id, api):
+    query = await api.users.get(user_ids=user_id, fields='sex')
+    return query[0]['sex'] if len(query) > 0 else 2

kurocore/utils/vk/keyboard.py

@@ -0,0 +1,184 @@
+import json
+from enum import Enum
+from json import JSONEncoder
+
+
+class EnumEncoder(JSONEncoder):
+    def default(self, obj):
+        return obj.value
+
+
+class VkKeyboardColor(Enum):
+    PRIMARY = 'primary'
+    SECONDARY = 'secondary'
+    NEGATIVE = 'negative'
+    POSITIVE = 'positive'
+
+
+class VkKeyboard:
+    def __init__(self, one_time=False, inline=False):
+        self.inline = inline
+        self.lines = [[]]
+
+        self.keyboard = {
+            'inline': self.inline,
+            'buttons': self.lines
+        }
+        if not inline:
+            self.keyboard['one_time'] = one_time
+
+    def __load_payload(self, payload) -> str:
+        if isinstance(payload, str):
+            return payload
+        elif isinstance(payload, dict):
+            return json.dumps(payload)
+        elif isinstance(payload, Payload):
+            return json.dumps(payload.get())
+
+    def add_text_button(self, text, payload=None, color=VkKeyboardColor.PRIMARY):
+        current_line = self.lines[-1]
+        if len(current_line) == 5:
+            raise TypeError('max elements in line: 5')
+
+        action = {
+            'type': 'text',
+            'label': text
+        }
+        if payload:
+            action.update({'payload': self.__load_payload(payload)})
+        button = {
+            'color': color,
+            'action': action
+        }
+        current_line.append(button)
+
+    def add_link_button(self, text, link, payload=None):
+        current_line = self.lines[-1]
+        if len(current_line) == 5:
+            raise TypeError('max elements in line: 5')
+
+        action = {
+            'type': 'open_link',
+            'link': link,
+            'label': text
+        }
+        if payload:
+            action.update({'payload': self.__load_payload(payload)})
+        button = {
+            'action': action
+        }
+        current_line.append(button)
+
+    def add_location_button(self, payload=None):
+        current_line = self.lines[-1]
+        if len(current_line) == 5:
+            raise TypeError('max elements in line: 5')
+
+        action = {
+            'type': 'location'
+        }
+
+        if payload:
+            action.update({'payload': self.__load_payload(payload)})
+
+        button = {
+            'action': action
+        }
+        current_line.append(button)
+
+    def add_vk_pay_button(self, hash):
+        current_line = self.lines[-1]
+        if len(current_line) == 5:
+            raise TypeError('max elements in line: 5')
+
+        action = {
+            'type': 'vkpay',
+            'hash': hash
+        }
+        button = {
+            'action': action
+        }
+        current_line.append(button)
+
+    def add_vk_apps_button(self, label, app_id, owner_id=None, hash=None, payload=None):
+        current_line = self.lines[-1]
+        if len(current_line) == 5:
+            raise TypeError('max elements in line: 5')
+
+        action = {
+            'type': 'open_app',
+            'label': label,
+            'app_id': app_id
+        }
+        if owner_id:
+            action.update({'owner_id': owner_id})
+        if hash:
+            action.update({'hash': hash})
+        if payload:
+            action.update({'payload': self.__load_payload(payload)})
+
+        button = {
+            'action': action
+        }
+        current_line.append(button)
+
+    def add_callback_button(self, label, payload=None, color=VkKeyboardColor.PRIMARY):
+        current_line = self.lines[-1]
+        if len(current_line) == 5:
+            raise TypeError('max elements in line: 5')
+
+        action = {
+            'type': 'callback',
+            'label': label
+        }
+
+        if payload:
+            action.update({'payload': self.__load_payload(payload)})
+
+        button = {
+            'action': action,
+            'color': color
+        }
+        current_line.append(button)
+
+    def add_line(self):
+        if len(self.lines) == 10:
+            if self.inline:
+                raise TypeError('max lines: 6')
+            else:
+                raise TypeError('max lines: 10')
+        self.lines.append([])
+
+    def get_current_line(self):
+        return self.lines[-1]
+
+    def get_keyboard(self):
+        keyboard = self.keyboard.copy()
+        if 'buttons' not in keyboard:
+            keyboard.update({'buttons': self.lines})
+        return json.dumps(keyboard, cls=EnumEncoder)
+
+    @classmethod
+    def get_empty_keyboard(cls):
+        keyboard = cls(True)
+        keyboard.keyboard['buttons'] = []
+        return keyboard.get_keyboard()
+
+    @classmethod
+    def get_empty_inline(cls):
+        keyboard = cls(False, True)
+        keyboard.keyboard['buttons'] = []
+        return keyboard.get_keyboard()
+
+
+class Payload:
+    def __init__(self, cmd, **kwargs):
+        self.cmd = cmd
+        self.args: dict = kwargs
+        self.value = {'command': cmd}
+
+    def get(self):
+        if self.args:
+            return {'command': self.cmd, 'args': self.args}
+        else:
+            return {'command': self.cmd}

kurocore/utils/vk/longpoll.py

@@ -0,0 +1,170 @@
+from enum import Enum
+from logging import Logger
+
+from kurocore import Config
+from kurocore.logger import RequestLogger
+from kurocore.utils.vk.vk import VK
+
+
+class DotDict(dict):
+    __getattr__ = dict.get
+    __setattr__ = dict.__setitem__
+    __delattr__ = dict.__delitem__
+
+
+class VKLongPoll:
+    __slots__ = ('vk', 'api', 'config', 'server', 'key', 'ts')
+
+    def __init__(self, config: Config, vk: VK):
+        self.vk = vk
+        self.api = vk.get_api()
+        self.config = config
+
+        self.server = ''
+        self.key = ''
+        self.ts = 0
+
+    async def init_lp(self):
+        group = (await self.api.groups.getById())['groups'][0]
+        lp = await self.api.groups.getLongPollServer(group_id=group['id'])
+
+        self.server = lp['server']
+        self.key = lp['key']
+        self.ts = lp['ts']
+
+    async def check(self):
+        async with self.vk.session.get(f'{self.server}?act=a_check&key={self.key}&ts={self.ts}&wait=25') as res:
+            body = await res.json()
+
+        if 'failed' in body:
+            code = body['failed']
+            if code == 1:
+                self.ts = body['ts']
+            if code == 2 or code == 3:
+                await self.init_lp()
+        else:
+            self.ts = body['ts']
+            
+            if self.config.logs.requests:
+                RequestLogger.log.info(body)
+
+            for event in body['updates']:
+                yield VkBotEvent(event)
+
+    async def listen(self):
+        while True:
+            async for event in self.check():
+                Logger.log.debug(f'new event: {event.raw}')
+                yield event
+
+
+class VkBotEvent(object):
+    __slots__ = (
+        'raw',
+        't', 'type',
+        'obj', 'object',
+        'client_info', 'message',
+        'group_id'
+    )
+
+    def __init__(self, raw):
+        self.raw = raw
+
+        try:
+            self.type = VkBotEventType(raw['type'])
+        except ValueError:
+            self.type = raw['type']
+
+        self.t = self.type  # shortcut
+
+        self.object = DotDict(raw['object'])
+        try:
+            self.message = DotDict(raw['object']['message'])
+        except KeyError:
+            self.message = None
+        self.obj = self.object
+        try:
+            self.client_info = DotDict(raw['object']['client_info'])
+        except KeyError:
+            self.client_info = None
+
+        self.group_id = raw['group_id']
+
+    def __repr__(self):
+        return '<{}({})>'.format(type(self), self.raw)
+
+
+class VkBotEventType(Enum):
+    MESSAGE_NEW = 'message_new'
+    MESSAGE_REPLY = 'message_reply'
+    MESSAGE_EDIT = 'message_edit'
+    MESSAGE_EVENT = 'message_event'
+
+    MESSAGE_TYPING_STATE = 'message_typing_state'
+
+    MESSAGE_ALLOW = 'message_allow'
+
+    MESSAGE_DENY = 'message_deny'
+
+    PHOTO_NEW = 'photo_new'
+
+    PHOTO_COMMENT_NEW = 'photo_comment_new'
+    PHOTO_COMMENT_EDIT = 'photo_comment_edit'
+    PHOTO_COMMENT_RESTORE = 'photo_comment_restore'
+
+    PHOTO_COMMENT_DELETE = 'photo_comment_delete'
+
+    AUDIO_NEW = 'audio_new'
+
+    VIDEO_NEW = 'video_new'
+
+    VIDEO_COMMENT_NEW = 'video_comment_new'
+    VIDEO_COMMENT_EDIT = 'video_comment_edit'
+    VIDEO_COMMENT_RESTORE = 'video_comment_restore'
+
+    VIDEO_COMMENT_DELETE = 'video_comment_delete'
+
+    WALL_POST_NEW = 'wall_post_new'
+    WALL_REPOST = 'wall_repost'
+
+    WALL_REPLY_NEW = 'wall_reply_new'
+    WALL_REPLY_EDIT = 'wall_reply_edit'
+    WALL_REPLY_RESTORE = 'wall_reply_restore'
+
+    WALL_REPLY_DELETE = 'wall_reply_delete'
+
+    BOARD_POST_NEW = 'board_post_new'
+    BOARD_POST_EDIT = 'board_post_edit'
+    BOARD_POST_RESTORE = 'board_post_restore'
+
+    BOARD_POST_DELETE = 'board_post_delete'
+
+    MARKET_COMMENT_NEW = 'market_comment_new'
+    MARKET_COMMENT_EDIT = 'market_comment_edit'
+    MARKET_COMMENT_RESTORE = 'market_comment_restore'
+
+    MARKET_COMMENT_DELETE = 'market_comment_delete'
+
+    GROUP_LEAVE = 'group_leave'
+    GROUP_JOIN = 'group_join'
+
+    USER_BLOCK = 'user_block'
+    USER_UNBLOCK = 'user_unblock'
+
+    POLL_VOTE_NEW = 'poll_vote_new'
+
+    GROUP_OFFICERS_EDIT = 'group_officers_edit'
+    GROUP_CHANGE_SETTINGS = 'group_change_settings'
+    GROUP_CHANGE_PHOTO = 'group_change_photo'
+
+    VKPAY_TRANSACTION = 'vkpay_transaction'
+
+    APP_PAYLOAD = 'app_payload'
+
+    DONUT_SUBSCRIPTION_CREATE = 'donut_subscription_create'
+    DONUT_SUBSCRIPTION_PROLONGED = 'donut_subscription_prolonged'
+    DONUT_SUBSCRIPTION_EXPIRED = 'donut_subscription_expired'
+    DONUT_SUBSCRIPTION_CANCELLED = 'donut_subscription_cancelled'
+    DONUT_SUBSCRIPTION_PRICE_CHANGED = 'donut_subscription_price_changed'
+    DONUT_SUBSCRIPTION_WITHDRAW = 'donut_money_withdraw'
+    DONUT_SUBSCRIPTION_WITHDRAW_ERROR = 'donut_money_withdraw_error'

kurocore/utils/vk/upload.py

@@ -0,0 +1,64 @@
+import json
+
+from aiohttp import FormData
+
+from handler.message import load_attachments
+from kurocore import LoggingClientSession
+from utils.vk.vk import VK, VkApiMethod
+
+
+class JsonParser:
+    @staticmethod
+    def dumps(data):
+        return json.dumps(data, ensure_ascii=False, separators=(",", ":"))
+
+    @staticmethod
+    def loads(string):
+        return json.loads(string)
+
+
+class VkUpload(object):
+    __slots__ = ('vk',)
+
+    def __init__(self, vk):
+        if not isinstance(vk, (VK, VkApiMethod)):
+            raise TypeError('The arg should be VK or VkApiMethod instance')
+
+        if isinstance(vk, VkApiMethod):
+            self.vk = vk
+        else:
+            self.vk = vk.get_api()
+
+    async def photo_messages(self, photo, pid=0):
+        upload_info = await self.vk.photos.getMessagesUploadServer(peer_id=pid)
+
+        data = FormData()
+        data.add_field(
+            'photo', photo,
+            content_type='multipart/form-data',
+            filename='a.png',
+        )
+
+        async with LoggingClientSession() as session, session.post(upload_info['upload_url'], data=data) as response:
+            response = await response.text()
+            response = json.loads(response)
+
+        photos = await self.vk.photos.saveMessagesPhoto(**response)
+        photos = [{'type': 'photo', 'photo': photo} for photo in photos]
+        return load_attachments(photos)
+
+    async def doc_message(self, doc, pid):
+        upload_info = await self.vk.docs.getMessagesUploadServer(peer_id=pid)
+
+        data = FormData()
+        data.add_field(
+            'file', doc,
+            content_type='multipart/form-data',
+            filename=f'a.png',
+        )
+
+        async with LoggingClientSession() as session, session.post(upload_info['upload_url'], data=data) as response:
+            response = await response.text()
+            response = json.loads(response)
+
+        return load_attachments([await self.vk.docs.save(**response)])

kurocore/utils/vk/utils.py

@@ -0,0 +1,6 @@
+from json import JSONEncoder
+
+
+class EnumEncoder(JSONEncoder):
+    def default(self, obj):
+        return obj.value

kurocore/utils/vk/vk.py

@@ -0,0 +1,117 @@
+import asyncio
+from time import time
+from typing import Union
+
+from aiohttp import ClientSession, FormData
+
+from kurocore.logger import RequestLogger
+
+
+class VkToken:
+    __slots__ = (
+        'cur_requests', 'max_rps',
+        '__token', '__last_req'
+    )
+
+    def __init__(self, token):
+        self.cur_requests = 0
+        self.max_rps = 20
+        self.__token = token
+        self.__last_req = 0
+
+    def __call__(self):
+        # if self.cur_requests >= self.max_rps:
+        #     raise TypeError('too many requests')
+        # self.cur_requests += 1
+        # self.__last_req = int(time())
+        # if self.__last_req < int(time()):
+        #     self.cur_requests = 0
+        return self.__token
+
+
+class VkTokenProvider:
+    __slots__ = (
+        'tokens',
+    )
+
+    def __init__(self, tokens: (Union[VkToken], VkToken)):
+        if type(tokens) is str:
+            self.tokens = [VkToken(tokens)]
+        else:
+            self.tokens = [VkToken(t) for t in tokens]
+
+    def obtain_token(self):
+        return self.tokens[0]()
+        # for t in self.tokens:
+        #     if t.cur_requests < t.max_rps:
+        #         return t()
+        # else:
+        #     raise ValueError('no free tokens!')
+
+
+class VK:
+    __slots__ = (
+        'token_provider',
+        'v', 'session'
+    )
+
+    def __init__(self, tokens: (str, list, tuple, set, frozenset), v='5.220'):
+        self.token_provider = VkTokenProvider(tokens)
+        self.v = v
+        self.session = ClientSession()
+
+    def shutdown(self):
+        asyncio.get_event_loop().run_until_complete(self.session.close())
+
+    async def call_method(self, method, **params):
+        params.update({'v': self.v})
+
+        params.update({'access_token': self.token_provider.obtain_token()})
+
+        async with self.session.post(
+                f'https://api.vk.com/method/{method}',
+                data=FormData(params)
+        ) as res:
+            j = await res.json()
+
+            if 'error' in j:
+                error = j['error']
+                raise VKApiException(error['error_msg'], error['error_code'])
+
+            params.update({"access_token": "[MASKED]"})
+            RequestLogger.log.debug(f'method: {method} {params}')
+
+            if 'response' in j:
+                return j['response']
+
+            return j
+
+    def get_api(self):
+        return VkApiMethod(self)
+
+
+class VkApiMethod(object):
+    __slots__ = ('_vk', '_method')
+
+    def __init__(self, vk, method=None):
+        self._vk: VK = vk
+        self._method = method
+
+    def __getattr__(self, method):
+        if '_' in method:
+            m = method.split('_')
+            method = m[0] + ''.join(i.title() for i in m[1:])
+
+        return VkApiMethod(
+            self._vk,
+            (self._method + '.' if self._method else '') + method
+        )
+
+    async def __call__(self, **kwargs):
+        return await self._vk.call_method(self._method, **kwargs)
+
+
+class VKApiException(Exception):
+    def __init__(self, msg, code):
+        self.msg = msg
+        self.code = code

kurocore/utils/vk_utils.py

@@ -0,0 +1,32 @@
+import io
+from random import randint
+
+from aiohttp import ClientSession
+
+
+def get_self_id(api):
+    return api.groups.getById()[0]['id']
+
+
+def generate_random_id():
+    return randint(-9 ** 99, 9 ** 99)
+
+
+async def get_user_name(id, api, name_case='nom'):
+    user = (await api.users.get(user_ids=id, name_case=name_case))[0]
+    return f'{user["first_name"]} {user["last_name"]}'
+
+
+async def reupload_attachments(attachments, upload):
+    new_attachments = []
+    for a in attachments:
+        t = a['type']
+        if t != 'photo':
+            continue
+        url = a[t]['sizes'][-1]['url']
+
+        async with ClientSession() as session, session.get(url) as response:
+            file = io.BytesIO(await response.content.read())
+            attachment = upload.photo_messages(file)[0]
+            new_attachments.append(f'photo{attachment["owner_id"]}_{attachment["id"]}')
+    return new_attachments