<commit_before>import asyncio
import datetime
import logging
import os
import sqlite3
from typing import Dict

from .. import api, crypto, utils
from ..params import KeeperParams
from ..proto import enterprise_pb2
from . import sqlite_storage, sox_data
from .storage_types import StorageRecord, StorageUser, StorageUserRecordLink, StorageTeam, \
    StorageRecordPermissions, StorageTeamUserLink, StorageSharedFolderRecordLink, StorageSharedFolderUserLink, \
    StorageSharedFolderTeamLink

API_SOX_REQUEST_USER_LIMIT = 1000


def get_prelim_data(params, enterprise_id=0, rebuild=False, min_updated=0, cache_only=False):
    # type: (KeeperParams, int, bool, int, bool) -> sox_data.SoxData
    def sync_down(name_by_id, store):  # type: (Dict[int, str], sqlite_storage.SqliteSoxStorage) ->  None
        def to_storage_types(user_data, username_lookup):
            def to_record_entity(record):
                entity = StorageRecord()
                entity.record_uid_bytes = record.recordUid
                entity.record_uid = utils.base64_url_encode(record.recordUid)
                entity.encrypted_data = record.encryptedData
                entity.shared = record.shared
                return entity

            def to_user_entity(user, email_lookup):
                entity = StorageUser()
                entity.status = user.status
                entity.user_uid = user.enterpriseUserId
                entity.email = email_lookup.get(entity.user_uid)
                return entity

            def to_user_record_link(uuid, ruid):
                link = StorageUserRecordLink()
                link.user_uid = uuid
                link.record_uid = ruid
                return link

            user_ent = to_user_entity(user_data, username_lookup)
            record_ents = [to_record_entity(x) for x in user_data.auditUserRecords if x.encryptedData]
            user_rec_links = {to_user_record_link(user_ent.user_uid, record_ent.record_uid) for record_ent in
                              record_ents}
            return user_ent, record_ents, user_rec_links

        def sync_all():
            print('Loading record information.', end='', flush=True)
            user_ids = list(user_lookup.keys())
            users, records, links = [], [], []
            while user_ids:
                print('.', end='', flush=True)
                token = b''
                chunk = user_ids[:API_SOX_REQUEST_USER_LIMIT]
                user_ids = user_ids[API_SOX_REQUEST_USER_LIMIT:]
                rq = enterprise_pb2.PreliminaryComplianceDataRequest()
                rq.enterpriseUserIds.extend(chunk)
                rq.includeNonShared = True
                has_more = True
                while has_more:
                    rq.continuationToken = token or rq.continuationToken
                    endpoint = 'enterprise/get_preliminary_compliance_data'
                    rs_type = enterprise_pb2.PreliminaryComplianceDataResponse
                    rs = api.communicate_rest(params, rq, endpoint, rs_type=rs_type)
                    print('.', end='', flush=True)
                    has_more = rs.hasMore
                    token = rs.continuationToken
                    for user_data in rs.auditUserData:
                        t_user, t_recs, t_links = to_storage_types(user_data, name_by_id)
                        users += [t_user]
                        records += t_recs
                        links += t_links

            store.rebuild_prelim_data(users, records, links)

        sync_all()
        print('.')

    enterprise_id = enterprise_id or next(((x['node_id'] >> 32) for x in params.enterprise['nodes']), 0)
    path = os.path.dirname(os.path.abspath(params.config_filename or '1'))
    database_name = os.path.join(path, f'sox_{enterprise_id}.db')
    tree_key = params.enterprise['unencrypted_tree_key']
    ecc_key = utils.base64_url_decode(params.enterprise['keys']['ecc_encrypted_private_key'])
    ecc_key = crypto.decrypt_aes_v2(ecc_key, tree_key)
    key = crypto.load_ec_private_key(ecc_key)
    storage = sqlite_storage.SqliteSoxStorage(
        get_connection=lambda: sqlite3.connect(database_name), owner=params.user, database_name=database_name
    )
    last_updated = storage.last_prelim_data_update
    refresh_data = rebuild or not last_updated or min_updated > last_updated
    if refresh_data and not cache_only:
        user_lookup = {x['enterprise_user_id']: x['username'] for x in params.enterprise['users']}
        sync_down(user_lookup, storage)
    return sox_data.SoxData(ec_private_key=key, storage=storage)


def get_compliance_data(params, node_id, enterprise_id=0, rebuild=False, min_updated=0, no_cache=False):
    def sync_down(sdata, node_uid, user_node_id_lookup):
        def run_sync_tasks():
            async def do_tasks():
                print('Loading compliance data.', end='', flush=True)
                users_uids = [int(uid) for uid in sdata.get_users()]
                record_uids_raw = [rec.record_uid_bytes for rec in sdata.get_records().values()]
                limit = asyncio.Semaphore(10)
                max_len = API_SOX_REQUEST_USER_LIMIT
                total_ruids = len(record_uids_raw)
                ruid_chunks = [record_uids_raw[x:x + max_len] for x in range(0, total_ruids, max_len)]
                tasks = [sync_chunk(chunk, users_uids, limit) for chunk in ruid_chunks]
                await asyncio.gather(*tasks)

            old_loop = asyncio.get_event_loop()
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            try:
                loop.run_until_complete(do_tasks())
                sdata.storage.set_compliance_data_updated()
                print('.')
            except KeyboardInterrupt:
                logging.info('SIGINT received: cancelling pending tasks')
            finally:
                pending = [task for task in asyncio.Task.all_tasks() if not task.done()]
                for task in pending:
                    task.cancel()
                asyncio.set_event_loop(old_loop)
                loop.run_until_complete(loop.shutdown_asyncgens())
                loop.stop()
                loop.close()

        async def sync_chunk(chunk, uuids, limit):
            print('.', end='', flush=True)
            rs = await fetch_response(raw_ruids=chunk, user_uids=uuids, limit=limit)
            print('.', end='', flush=True)
            await save_response(rs)
            print(':', end='', flush=True)

        async def fetch_response(raw_ruids, user_uids, limit=None):
            async with limit:
                rq = enterprise_pb2.ComplianceReportRequest()
                rq.saveReport = False
                rq.reportName = f'Compliance Report on {datetime.datetime.now()}'
                report_run = rq.complianceReportRun
                report_run.users.extend(user_uids)
                report_run.records.extend(raw_ruids)
                caf = report_run.reportCriteriaAndFilter
                caf.nodeId = node_uid
                endpoint = 'enterprise/run_compliance_report'
                return api.communicate_rest(params, rq, endpoint, rs_type=enterprise_pb2.ComplianceReportResponse)

        anon_id = 0

        async def save_response(rs):
            def hash_anon_ids(response):
                # create new user uid for each anonymous user (uid >> 32 == 0)
                anon_ids = dict()
                nonlocal anon_id<commit_msg>fixed asyncio exception-handling for all supported python versions (#829)<commit_after>import asyncio
import datetime
import logging
import os
import sqlite3
from typing import Dict

from .. import api, crypto, utils
from ..params import KeeperParams
from ..proto import enterprise_pb2
from . import sqlite_storage, sox_data
from .storage_types import StorageRecord, StorageUser, StorageUserRecordLink, StorageTeam, \
    StorageRecordPermissions, StorageTeamUserLink, StorageSharedFolderRecordLink, StorageSharedFolderUserLink, \
    StorageSharedFolderTeamLink

API_SOX_REQUEST_USER_LIMIT = 1000


def get_prelim_data(params, enterprise_id=0, rebuild=False, min_updated=0, cache_only=False):
    # type: (KeeperParams, int, bool, int, bool) -> sox_data.SoxData
    def sync_down(name_by_id, store):  # type: (Dict[int, str], sqlite_storage.SqliteSoxStorage) ->  None
        def to_storage_types(user_data, username_lookup):
            def to_record_entity(record):
                entity = StorageRecord()
                entity.record_uid_bytes = record.recordUid
                entity.record_uid = utils.base64_url_encode(record.recordUid)
                entity.encrypted_data = record.encryptedData
                entity.shared = record.shared
                return entity

            def to_user_entity(user, email_lookup):
                entity = StorageUser()
                entity.status = user.status
                entity.user_uid = user.enterpriseUserId
                entity.email = email_lookup.get(entity.user_uid)
                return entity

            def to_user_record_link(uuid, ruid):
                link = StorageUserRecordLink()
                link.user_uid = uuid
                link.record_uid = ruid
                return link

            user_ent = to_user_entity(user_data, username_lookup)
            record_ents = [to_record_entity(x) for x in user_data.auditUserRecords if x.encryptedData]
            user_rec_links = {to_user_record_link(user_ent.user_uid, record_ent.record_uid) for record_ent in
                              record_ents}
            return user_ent, record_ents, user_rec_links

        def sync_all():
            print('Loading record information.', end='', flush=True)
            user_ids = list(user_lookup.keys())
            users, records, links = [], [], []
            while user_ids:
                print('.', end='', flush=True)
                token = b''
                chunk = user_ids[:API_SOX_REQUEST_USER_LIMIT]
                user_ids = user_ids[API_SOX_REQUEST_USER_LIMIT:]
                rq = enterprise_pb2.PreliminaryComplianceDataRequest()
                rq.enterpriseUserIds.extend(chunk)
                rq.includeNonShared = True
                has_more = True
                while has_more:
                    rq.continuationToken = token or rq.continuationToken
                    endpoint = 'enterprise/get_preliminary_compliance_data'
                    rs_type = enterprise_pb2.PreliminaryComplianceDataResponse
                    rs = api.communicate_rest(params, rq, endpoint, rs_type=rs_type)
                    print('.', end='', flush=True)
                    has_more = rs.hasMore
                    token = rs.continuationToken
                    for user_data in rs.auditUserData:
                        t_user, t_recs, t_links = to_storage_types(user_data, name_by_id)
                        users += [t_user]
                        records += t_recs
                        links += t_links

            store.rebuild_prelim_data(users, records, links)

        sync_all()
        print('.')

    enterprise_id = enterprise_id or next(((x['node_id'] >> 32) for x in params.enterprise['nodes']), 0)
    path = os.path.dirname(os.path.abspath(params.config_filename or '1'))
    database_name = os.path.join(path, f'sox_{enterprise_id}.db')
    tree_key = params.enterprise['unencrypted_tree_key']
    ecc_key = utils.base64_url_decode(params.enterprise['keys']['ecc_encrypted_private_key'])
    ecc_key = crypto.decrypt_aes_v2(ecc_key, tree_key)
    key = crypto.load_ec_private_key(ecc_key)
    storage = sqlite_storage.SqliteSoxStorage(
        get_connection=lambda: sqlite3.connect(database_name), owner=params.user, database_name=database_name
    )
    last_updated = storage.last_prelim_data_update
    refresh_data = rebuild or not last_updated or min_updated > last_updated
    if refresh_data and not cache_only:
        user_lookup = {x['enterprise_user_id']: x['username'] for x in params.enterprise['users']}
        sync_down(user_lookup, storage)
    return sox_data.SoxData(ec_private_key=key, storage=storage)


def get_compliance_data(params, node_id, enterprise_id=0, rebuild=False, min_updated=0, no_cache=False):
    def sync_down(sdata, node_uid, user_node_id_lookup):
        def run_sync_tasks():
            async def do_tasks(return_exceptions=False):
                print('Loading compliance data.', end='', flush=True)
                users_uids = [int(uid) for uid in sdata.get_users()]
                record_uids_raw = [rec.record_uid_bytes for rec in sdata.get_records().values()]
                limit = asyncio.Semaphore(10)
                max_len = API_SOX_REQUEST_USER_LIMIT
                total_ruids = len(record_uids_raw)
                ruid_chunks = [record_uids_raw[x:x + max_len] for x in range(0, total_ruids, max_len)]
                tasks = [sync_chunk(chunk, users_uids, limit) for chunk in ruid_chunks]
                await asyncio.gather(*tasks, return_exceptions=return_exceptions)
                sdata.storage.set_compliance_data_updated()

            py_version_3_6 = not hasattr(asyncio, 'run')
            if not py_version_3_6:
                try:
                    asyncio.run(do_tasks(True))
                finally:
                    print('')
            else:
                old_loop = asyncio.get_event_loop()
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
                try:
                    loop.run_until_complete(do_tasks())
                except KeyboardInterrupt:
                    logging.info('SIGINT received: cancelling pending tasks')
                finally:
                    try:
                        pending = [task for task in asyncio.Task.all_tasks() if not task.done()]
                        for task in pending:
                            task.cancel()
                        loop.run_until_complete(loop.shutdown_asyncgens())
                    finally:
                        asyncio.set_event_loop(old_loop)
                        loop.stop()
                        loop.close()

        async def sync_chunk(chunk, uuids, limit):
            print('.', end='', flush=True)
            rs = await fetch_response(raw_ruids=chunk, user_uids=uuids, limit=limit)
            print('.', end='', flush=True)
            await save_response(rs)
            print(':', end='', flush=True)

        async def fetch_response(raw_ruids, user_uids, limit=None):
            async with limit:
                rq = enterprise_pb2.ComplianceReportRequest()
                rq.saveReport = False
                rq.reportName = f'Compliance Report on {datetime.datetime.now()}'
                report_run = rq.complianceReportRun
                report_run.users.extend(user_uids)
                report_run.records.extend(raw_ruids)
                caf = report_run.reportCriteriaAndFilter
                caf.nodeId = node_uid
                endpoint = 'enterprise/run_compliance_report'
                return api.communicate_rest(params, rq, endpoint, rs_type=enterprise_pb2.ComplianceReportResponse)

        anon_id = 0

        async def save_response(rs):
            def hash_anon_ids(response):
                # create new user uid for each anonymous user (uid >> 32 == 0)
                anon_ids = dict()
                nonlocal anon_id
{'fraction_non_alphanumeric': 0.0843890079906451, 'fraction_numerical': 0.004482557006431495, 'mean_word_length': 2.142273933455807, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '221971', 'n_tokens_mistral': 4505, 'n_tokens_neox': 4271, 'n_words': 877}