Merge pull request #22 from soxoj/self-check-fix

Self-checking mode fixed, tags/names site filtering & ranking

maigret/maigret.py

@@ -368,7 +368,6 @@ async def maigret(username, site_dict, query_notify, logger,
         results_site['parsing_enabled'] = recursive_search
         results_site['url_main'] = site.url_main
 
-
         headers = {
             'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 11.1; rv:55.0) Gecko/20100101 Firefox/55.0',
         }
@@ -506,95 +505,102 @@ def timeout_check(value):
     return timeout
 
 
-async def site_self_check(site_name, site_data, logger, no_progressbar=False):
+async def site_self_check(site, logger, semaphore, db: MaigretDatabase, no_progressbar=False):
     query_notify = Mock()
     changes = {
         'disabled': False,
     }
 
-    check_data = [
-        (site_data.username_claimed, QueryStatus.CLAIMED),
-        (site_data.username_unclaimed, QueryStatus.AVAILABLE),
-    ]
+    try:
+        check_data = [
+            (site.username_claimed, QueryStatus.CLAIMED),
+            (site.username_unclaimed, QueryStatus.AVAILABLE),
+        ]
+    except:
+        print(site.__dict__)
 
-    logger.info(f'Checking {site_name}...')
+    logger.info(f'Checking {site.name}...')
 
     for username, status in check_data:
-        results = await maigret(
-            username,
-            {site_name: site_data},
-            query_notify,
-            logger,
-            timeout=30,
-            forced=True,
-            no_progressbar=no_progressbar,
-        )
+        async with semaphore:
+            results_dict = await maigret(
+                username,
+                {site.name: site},
+                query_notify,
+                logger,
+                timeout=30,
+                forced=True,
+                no_progressbar=True,
+            )
 
-        # don't disable entries with other ids types
-        if site_name not in results:
-            logger.info(results)
-            changes['disabled'] = True
-            continue
+            # don't disable entries with other ids types
+            # TODO: make normal checking
+            if site.name not in results_dict:
+                logger.info(results_dict)
+                changes['disabled'] = True
+                continue
+
+            result = results_dict[site.name]['status']
+
+
+        site_status = result.status
 
-        site_status = results[site_name]['status'].status
         if site_status != status:
             if site_status == QueryStatus.UNKNOWN:
-                msgs = site_data.absence_strs
-                etype = site_data.check_type
-                logger.info(f'Error while searching {username} in {site_name}: {msgs}, type {etype}')
+                msgs = site.absence_strs
+                etype = site.check_type
+                logger.warning(f'Error while searching {username} in {site.name}: {result.context}, {msgs}, type {etype}')
                 # don't disable in case of available username
                 if status == QueryStatus.CLAIMED:
                     changes['disabled'] = True
             elif status == QueryStatus.CLAIMED:
-                logger.info(f'Not found `{username}` in {site_name}, must be claimed')
-                logger.info(results[site_name])
+                logger.warning(f'Not found `{username}` in {site.name}, must be claimed')
+                logger.info(results_dict[site.name])
                 changes['disabled'] = True
             else:
-                logger.info(f'Found `{username}` in {site_name}, must be available')
-                logger.info(results[site_name])
+                logger.warning(f'Found `{username}` in {site.name}, must be available')
+                logger.info(results_dict[site.name])
                 changes['disabled'] = True
 
-    logger.info(f'Site {site_name} checking is finished')
+    logger.info(f'Site {site.name} checking is finished')
+
+    if changes['disabled'] != site.disabled:
+        site.disabled = changes['disabled']
+        db.update_site(site)
+        action = 'Disabled' if not site.disabled else 'Enabled'
+        print(f'{action} site {site.name}...')
+
     return changes
 
 
-async def self_check(json_file, logger):
-    db = MaigretDatabase()
-    db.load_from_file(json_file)
-    sites = db.sites
-    all_sites = {}
+async def self_check(db: MaigretDatabase, site_data: dict, logger):
+    sem = asyncio.Semaphore(10)
+    tasks = []
+    all_sites = site_data
 
-    def disabled_count(data):
-        return len(list(filter(lambda x: x.get('disabled', False), data)))
-
-    async def update_site_data(site_name, site_data, all_sites, logger):
-        updates = await site_self_check(site_name, dict(site_data), logger)
-        all_sites[site_name].update(updates)
-
-    for site in sites:
-        all_sites[site.name] = site.information
+    def disabled_count(lst):
+        return len(list(filter(lambda x: x.disabled, lst)))
 
     disabled_old_count = disabled_count(all_sites.values())
 
-    tasks = []
-    for site_name, site_data in all_sites.items():
-        future = asyncio.ensure_future(update_site_data(site_name, site_data, all_sites, logger))
+    for _, site in all_sites.items():
+        check_coro = site_self_check(site, logger, sem, db)
+        future = asyncio.ensure_future(check_coro)
         tasks.append(future)
 
-    await asyncio.gather(*tasks)
+    for f in tqdm.asyncio.tqdm.as_completed(tasks):
+        await f
 
     disabled_new_count = disabled_count(all_sites.values())
     total_disabled = disabled_new_count - disabled_old_count
-    if total_disabled > 0:
+
+    if total_disabled >= 0:
         message = 'Disabled'
     else:
         message = 'Enabled'
         total_disabled *= -1
-    print(f'{message} {total_disabled} checked sites. Run with `--info` flag to get more information')
 
-    with open(json_file, 'w') as f:
-        data['sites'] = all_sites
-        json.dump(data, f, indent=4)
+    print(f'{message} {total_disabled} checked sites. Run with `--info` flag to get more information')
 
 
 async def main():
@@ -621,9 +627,6 @@ async def main():
                         action="store_true", dest="debug", default=False,
                         help="Saving debugging information and sites responses in debug.txt."
                         )
-    parser.add_argument("--rank", "-r",
-                        action="store_true", dest="rank", default=False,
-                        help="Present websites ordered by their Alexa.com global rank in popularity.")
     parser.add_argument("--folderoutput", "-fo", dest="folderoutput", default="reports",
                         help="If using multiple usernames, the output of the results will be saved to this folder."
                         )
@@ -637,7 +640,7 @@ async def main():
                         )
     parser.add_argument("--site",
                         action="append", metavar='SITE_NAME',
-                        dest="site_list", default=None,
+                        dest="site_list", default=[],
                         help="Limit analysis to just the listed sites (use several times to specify more than one)"
                         )
     parser.add_argument("--proxy", "-p", metavar='PROXY_URL',
@@ -758,7 +761,7 @@ async def main():
                 usernames[v] = k
 
     if args.tags:
-        args.tags = set(str(args.tags).split(','))
+        args.tags = list(set(str(args.tags).split(',')))
 
     if args.json_file is None:
         args.json_file = \
@@ -766,53 +769,40 @@ async def main():
                          "resources/data.json"
                          )
 
-    # Database self-checking
-    if args.self_check:
-        print('Maigret sites database self-checking...')
-        await self_check(args.json_file, logger)
+    if args.top_sites == 0:
+        args.top_sites = sys.maxsize
 
     # Create object with all information about sites we are aware of.
     try:
         db = MaigretDatabase().load_from_file(args.json_file)
-        site_data_all = db.ranked_sites_dict(top=args.top_sites)
+        site_data = db.ranked_sites_dict(top=args.top_sites, tags=args.tags, names=args.site_list)
     except Exception as error:
         print(f"ERROR:  {error}")
         sys.exit(1)
 
-    if args.site_list is None:
-        # Not desired to look at a sub-set of sites
-        site_data = site_data_all
-    else:
-        # User desires to selectively run queries on a sub-set of the site list.
-
-        # Make sure that the sites are supported & build up pruned site database.
-        site_data = {}
-        site_missing = []
-        for site in args.site_list:
-            for existing_site in site_data_all:
-                if site.lower() == existing_site.lower():
-                    site_data[existing_site] = site_data_all[existing_site]
-            if not site_data:
-                # Build up list of sites not supported for future error message.
-                site_missing.append(f"'{site}'")
-
-        if site_missing:
-            print(
-                f"Error: Desired sites not found: {', '.join(site_missing)}.")
-            sys.exit(1)
-
-    if args.rank:
-        # Sort data by rank
-        site_dataCpy = dict(site_data)
-        ranked_sites = sorted(site_data, key=lambda k: ("rank" not in k, site_data[k].get("rank", sys.maxsize)))
-        site_data = {}
-        for site in ranked_sites:
-            site_data[site] = site_dataCpy.get(site)
+    # Database self-checking
+    if args.self_check:
+        print('Maigret sites database self-checking...')
+        await self_check(db, site_data, logger)
+        if input('Do you want to save changes permanently? [yYnN]\n').lower() == 'y':
+            db.save_to_file(args.json_file)
+            print('Database was successfully updated.')
+        else:
+            print('Updates will be applied only for current search session.')
 
     # Database consistency
     enabled_count = len(list(filter(lambda x: not x.disabled, site_data.values())))
     print(f'Sites in database, enabled/total: {enabled_count}/{len(site_data)}')
 
+    if not enabled_count:
+        print('No sites to check, exiting!')
+        sys.exit(2)
+
+    if usernames == ['-']:
+        # magic params to exit after init
+        print('No usernames to check, exiting.')
+        sys.exit(0)
+
     # Create notify object for query results.
     query_notify = QueryNotifyPrint(result=None,
                                     verbose=args.verbose,

maigret/sites.py

@@ -128,11 +128,18 @@ class MaigretDatabase:
     def sites_dict(self):
         return {site.name: site for site in self._sites}
 
-    def ranked_sites_dict(self, reverse=False, top=sys.maxsize, tags=[]):
-        if not tags:
+    def ranked_sites_dict(self, reverse=False, top=sys.maxsize, tags=[], names=[]):
+        normalized_names = list(map(str.lower, names))
+        normalized_tags = list(map(str.lower, tags))
+
+        is_tags_ok = lambda x: set(x.tags).intersection(set(normalized_tags))
+        is_name_ok = lambda x: x.name.lower() in normalized_names
+        is_engine_ok = lambda x: isinstance(x.engine, str) and x.engine.lower() in normalized_tags
+
+        if not tags and not names:
             filtered_list = self.sites
         else:
-            filtered_list = [s for s in self.sites if set(s.tags).intersection(set(tags)) or s.engine in tags]
+            filtered_list = [s for s in self.sites if is_tags_ok(s) or is_name_ok(s) or is_engine_ok(s)]
 
         sorted_list = sorted(filtered_list, key=lambda x: x.alexa_rank, reverse=reverse)[:top]
         return {site.name: site for site in sorted_list}

pytest.ini

@@ -0,0 +1,5 @@
+# pytest.ini
+[pytest]
+filterwarnings =
+    error
+    ignore::UserWarning

tests/conftest.py

@@ -0,0 +1,12 @@
+from _pytest.mark import Mark
+
+
+empty_mark = Mark('', [], {})
+
+
+def by_slow_marker(item):
+    return item.get_closest_marker('slow', default=empty_mark)
+
+
+def pytest_collection_modifyitems(items):
+    items.sort(key=by_slow_marker, reverse=False)

tests/test_maigret.py

@@ -0,0 +1,106 @@
+"""Maigret main module test functions"""
+import asyncio
+from mock import Mock
+import pytest
+
+from maigret.sites import MaigretDatabase, MaigretSite
+from maigret.maigret import self_check
+
+
+EXAMPLE_DB = {
+    'engines': {
+    },
+    'sites': {
+        "GooglePlayStore": {
+            "tags": [
+                "global",
+                "us"
+            ],
+            "disabled": False,
+            "checkType": "status_code",
+            "alexaRank": 1,
+            "url": "https://play.google.com/store/apps/developer?id={username}",
+            "urlMain": "https://play.google.com/store",
+            "usernameClaimed": "Facebook_nosuchname",
+            "usernameUnclaimed": "noonewouldeverusethis7"
+        },
+        "Reddit": {
+            "tags": [
+                "news",
+                "social",
+                "us"
+            ],
+            "checkType": "status_code",
+            "presenseStrs": [
+                "totalKarma"
+            ],
+            "disabled": True,
+            "alexaRank": 17,
+            "url": "https://www.reddit.com/user/{username}",
+            "urlMain": "https://www.reddit.com/",
+            "usernameClaimed": "blue",
+            "usernameUnclaimed": "noonewouldeverusethis7"
+        },
+    }
+}
+
+
+@pytest.mark.slow
+def test_self_check_db_positive_disable():
+    logger = Mock()
+    db = MaigretDatabase()
+    db.load_from_json(EXAMPLE_DB)
+
+    assert db.sites[0].disabled == False
+
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(self_check(db, db.sites_dict, logger))
+
+    assert db.sites[0].disabled == True
+
+
+@pytest.mark.slow
+def test_self_check_db_positive_enable():
+    logger = Mock()
+    db = MaigretDatabase()
+    db.load_from_json(EXAMPLE_DB)
+
+    db.sites[0].disabled = True
+    db.sites[0].username_claimed = 'Facebook'
+    assert db.sites[0].disabled == True
+
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(self_check(db, db.sites_dict, logger))
+
+    assert db.sites[0].disabled == False
+
+
+@pytest.mark.slow
+def test_self_check_db_negative_disabled():
+    logger = Mock()
+    db = MaigretDatabase()
+    db.load_from_json(EXAMPLE_DB)
+
+    db.sites[0].disabled = True
+    assert db.sites[0].disabled == True
+
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(self_check(db, db.sites_dict, logger))
+
+    assert db.sites[0].disabled == True
+
+
+@pytest.mark.slow
+def test_self_check_db_negative_enabled():
+    logger = Mock()
+    db = MaigretDatabase()
+    db.load_from_json(EXAMPLE_DB)
+
+    db.sites[0].disabled = False
+    db.sites[0].username_claimed = 'Facebook'
+    assert db.sites[0].disabled == False
+
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(self_check(db, db.sites_dict, logger))
+
+    assert db.sites[0].disabled == False

tests/test_sites.py

@@ -118,3 +118,14 @@ def test_ranked_sites_dict():
 
     # filtering by engine
     assert list(db.ranked_sites_dict(tags=['ucoz']).keys()) == ['3']
+
+    # filtering by names
+    assert list(db.ranked_sites_dict(names=['1', '2']).keys()) == ['1', '2']
+    assert list(db.ranked_sites_dict(names=['2', '3']).keys()) == ['2', '3']
+
+    # disjunction
+    assert list(db.ranked_sites_dict(names=['2'], tags=['forum']).keys()) == ['1', '2']
+    assert list(db.ranked_sites_dict(names=['2'], tags=['forum'], reverse=True).keys()) == ['2', '1']
+    assert list(db.ranked_sites_dict(names=['2'], tags=['ucoz']).keys()) == ['2', '3']
+    assert list(db.ranked_sites_dict(names=['4'], tags=['ru']).keys()) == ['2']
+    assert list(db.ranked_sites_dict(names=['4'], tags=['nosuchtag']).keys()) == []

utils/update_site_data.py

@@ -12,6 +12,8 @@ import xml.etree.ElementTree as ET
 from datetime import datetime
 from argparse import ArgumentParser, RawDescriptionHelpFormatter
 
+from maigret.maigret import MaigretDatabase
+
 RANKS = {str(i):str(i) for i in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 50, 100, 500]}
 RANKS.update({
     '1000': '1K',
@@ -22,7 +24,7 @@ RANKS.update({
     '50000000': '10M',
 })
 
-def get_rank(domain_to_query, dest, print_errors=True):
+def get_rank(domain_to_query, site, print_errors=True):
     #Retrieve ranking data via alexa API
     url = f"http://data.alexa.com/data?cli=10&url={domain_to_query}"
     xml_data = requests.get(url).text
@@ -30,16 +32,16 @@ def get_rank(domain_to_query, dest, print_errors=True):
 
     try:
         #Get ranking for this site.
-        dest['rank'] = int(root.find('.//REACH').attrib['RANK'])
+        site.alexa_rank = int(root.find('.//REACH').attrib['RANK'])
         country = root.find('.//COUNTRY')
         if not country is None and country.attrib:
             country_code = country.attrib['CODE']
-            tags = set(dest.get('tags', []))
+            tags = set(site.tags)
             if country_code:
                 tags.add(country_code.lower())
-            dest['tags'] = sorted(list(tags))
-            if 'type' in dest and dest['type'] != 'username':
-                dest['disabled'] = False
+            site.tags = sorted(list(tags))
+            if site.type != 'username':
+                site.disabled = False
     except Exception as e:
         if print_errors:
             logging.error(e)
@@ -67,38 +69,40 @@ if __name__ == '__main__':
                         dest="base_file", default="maigret/resources/data.json",
                         help="JSON file with sites data to update.")
 
+    parser.add_argument('--empty-only', help='update only sites without rating', action='store_true')
+
     pool = list()
 
     args = parser.parse_args()
 
-    with open(args.base_file, "r", encoding="utf-8") as data_file:
-        sites_info = json.load(data_file)
-        data = sites_info['sites']
-        engines = sites_info['engines']
+    db = MaigretDatabase()
+    sites_subset = db.load_from_file(args.base_file).sites
 
     with open("sites.md", "w") as site_file:
-        data_length = len(data)
         site_file.write(f"""
-## List of supported sites: total {data_length}\n
+## List of supported sites: total {len(sites_subset)}\n
 Rank data fetched from Alexa by domains.
 
 """)
 
-        for social_network in data:
-            url_main = data.get(social_network).get("urlMain")
-            data.get(social_network)["rank"] = 0
-            th = threading.Thread(target=get_rank, args=(url_main, data.get(social_network)))
-            pool.append((social_network, url_main, th))
+        for site in sites_subset:
+            url_main = site.url_main
+            if site.alexa_rank < sys.maxsize and args.empty_only:
+                continue
+            site.alexa_rank = 0
+            th = threading.Thread(target=get_rank, args=(url_main, site))
+            pool.append((site.name, url_main, th))
             th.start()
 
         index = 1
-        for social_network, url_main, th in pool:
+        for site_name, url_main, th in pool:
             th.join()
-            sys.stdout.write("\r{0}".format(f"Updated {index} out of {data_length} entries"))
+            sys.stdout.write("\r{0}".format(f"Updated {index} out of {len(sites_subset)} entries"))
             sys.stdout.flush()
             index = index + 1
 
-        sites_full_list = [(site, site_data['rank']) for site, site_data in data.items()]
+        sites_full_list = [(s, s.alexa_rank) for s in sites_subset]
+
         sites_full_list.sort(reverse=False, key=lambda x: x[1])
 
         while sites_full_list[0][1] == 0:
@@ -107,20 +111,17 @@ Rank data fetched from Alexa by domains.
 
         for num, site_tuple in enumerate(sites_full_list):
             site, rank = site_tuple
-            url_main = data[site]['urlMain']
+            url_main = site.url_main
             valid_rank = get_step_rank(rank)
-            all_tags = data[site].get('tags', [])
+            all_tags = site.tags
             tags = ', ' + ', '.join(all_tags) if all_tags else ''
             note = ''
-            if data[site].get('disabled'):
+            if site.disabled:
                 note = ', search is disabled'
             site_file.write(f'1. [{site}]({url_main})*: top {valid_rank}{tags}*{note}\n')
+            db.update_site(site)
 
         site_file.write(f'\nAlexa.com rank data fetched at ({datetime.utcnow()} UTC)\n')
-
-    sorted_json_data = json.dumps({'sites': data, 'engines': engines}, indent=2, sort_keys=True)
-
-    with open(args.base_file, "w") as data_file:
-        data_file.write(sorted_json_data)
+        db.save_to_file(args.base_file)
 
     print("\nFinished updating supported site listing!")