rok-stat/rok-stat-grabber.py

import argparse
import re
import io
import os
import sys
import time
from datetime import datetime
from ppadb.client import Client as AdbClient
from PIL import Image
import pytesseract

# ----
# Program
# ----

PROGRAM_NAME = "RoK Scraper"
PROGRAM_VERSION = "0.3"
PROGRAM_DESCRIPTION = "This program controls devices via ADB to take screenshots in Rise of Kingdoms."

# TODO:
# - Skip your own ranking
# - Correct for device rotation
# - Better logic for tap_profile()
# - Draw box around skipped profiles
# - Support for the last ~995-1000
# - Coordinates to percentages

# ----
# Arguments
# ----

parser = argparse.ArgumentParser(description = PROGRAM_DESCRIPTION)
parser.add_argument("--power", help = "get full power rankings", default = False, action = "store_true")
parser.add_argument("--top", help = "get stats on the top 300", default = False, action = "store_true")
parser.add_argument("-p", "--project", help = "project name", required = True)
parser.add_argument("-s", "--serial", help = "device serial", required = True)
parser.add_argument("-v", "--verbose", help = "be verbose", default = False, action = "store_true")
parser.add_argument("-c", "--count", help = "how many to scrape", default = 0, type = int)
parser.add_argument("-n", "--start", help = "what number to start scraping at", default = 1, type = int)
parser.add_argument("-k", "--skipped", help = "how many have been skipped", type = int)
parser.add_argument("--debug", help = "debug", default = False, action = "store_true")
arguments = parser.parse_args()

# ----
# Coordinates
# ----

ROTATE = 90 # TODO: fix this; shouldn't be needed

TAP_LOCATIONS = {
    "1": (460, 460),
    "2": (460, 624),
    "3": (460, 780),
    "4": (460, 940),
    "Middle": (460, 980),
    "Middle + 1": (460, 1136),
    "Middle + 2": (460, 1300)
}

# ----
# Functions
# ----

# TODO: pass device instead of image
def read_string_from_image(image, x, y, x2, y2):
    return pytesseract.image_to_string(image.crop((x, y, x2, y2)), config="--psm 6").strip().replace('\n', ' ').replace('\r', '').replace('\t', ' ')

def get_clipboard(device):
    raw = device.shell('am broadcast -n "ch.pete.adbclipboard/.ReadReceiver"')
    dataMatcher = re.compile("^.*\n.*data=\"(.*)\"$", re.DOTALL)
    dataMatch = dataMatcher.match(raw)
    return dataMatch.group(1)

def take_screenshot(device, path):
    screencap = device.screencap()
    image = Image.open(io.BytesIO(screencap)).rotate(ROTATE, expand=1)
    image.save(path + ".png", "PNG")
    return

# TODO: this can be done better
def get_datetime_string():
    d = datetime.now()
    return str(d.year).zfill(4) + str(d.month).zfill(2) + str(d.day).zfill(2) + str(d.hour).zfill(2) + str(d.minute).zfill(2) + str(d.second).zfill(2)

def read_single_power_ranking(device, folder):
    datetime = get_datetime_string()

    # Screenshot profile
    # ID, Username, Power, Kill Points
    time.sleep(1)
    take_screenshot(device, folder + datetime + "_profile")
    
    # Copy and save username
    device.shell("input tap 1055 455")
    time.sleep(1)
    username = get_clipboard(device)
    with open(projectFolder + "usernames.txt", 'a+', newline='', encoding='utf-8') as output_file:
        output_file.write(datetime + "\t")
        output_file.write(username)
        output_file.write("\n")
    if arguments.verbose: print(username, end = "", flush = True)
    
    # Click "More Info"
    # Screenshot "More Info"
    # Power, Kill Points, Highest Power, Victory, Defeat, Dead, Scout Times, Resources Gathered, Resource Assistance, Alliance Help Times
    device.shell("input tap 620 1070")
    time.sleep(2)
    take_screenshot(device, folder + datetime + "_more")
    
    # Click "?"
    # Screenshot Kills
    # Kill Points, T1-5 Kills, T1-5 Points
    device.shell("input tap 1725 250")
    time.sleep(1)
    take_screenshot(device, folder + datetime + "_kills")
    
    # Exit
    if arguments.verbose: print("") # TODO: say "done" or something here if successful
    device.shell("input tap 2230 85")
    time.sleep(1)
    device.shell("input tap 2185 160")
    return

# Tap to enter a profile, skipping ones that cannot be clicked
def tap_profile(i, folder): # TODO: fix this absolute mess of a function
    if i == 1:
        if is_profile(TAP_LOCATIONS["1"], folder, i): return 1
        elif is_profile(TAP_LOCATIONS["2"], folder, i + 1): return 2
        elif is_profile(TAP_LOCATIONS["3"], folder, i + 2): return 3
        elif is_profile(TAP_LOCATIONS["4"], folder, i + 3): return 4
        else: sys.exit("Profile not found.")
    elif i == 2:
        if is_profile(TAP_LOCATIONS["2"], folder, i): return 1
        elif is_profile(TAP_LOCATIONS["3"], folder, i + 1): return 2
        elif is_profile(TAP_LOCATIONS["4"], folder, i + 2): return 3
        else: sys.exit("Profile not found.")
    elif i == 3:
        if is_profile(TAP_LOCATIONS["3"], folder, i): return 1
        elif is_profile(TAP_LOCATIONS["4"], folder, i + 1): return 2
        else: sys.exit("Profile not found.")
    elif i == 4:
        if is_profile(TAP_LOCATIONS["4"], folder, i): return 1
        else: sys.exit("Profile not found.")
    else:
        if is_profile(TAP_LOCATIONS["Middle"], folder, i): return 1
        elif is_profile(TAP_LOCATIONS["Middle + 1"], folder, i + 1): return 2
        elif is_profile(TAP_LOCATIONS["Middle + 2"], folder, i + 2): return 3
        else: sys.exit("Profile not found.")

# Tap and test if you are in a profile 
def is_profile(coordinates, folder, i):
    device.shell("input tap " + str(coordinates[0]) + " " + str(coordinates[1]))
    time.sleep(1)
    if read_string_from_image(Image.open(io.BytesIO(device.screencap())).rotate(ROTATE, expand=1), 1652, 1167, 1745, 1209) == "Mail":
        return True
    else:
        device.shell("input tap " + str(coordinates[0]) + " " + str(coordinates[1]))
        time.sleep(2)
        if read_string_from_image(Image.open(io.BytesIO(device.screencap())).rotate(ROTATE, expand=1), 1652, 1167, 1745, 1209) == "Mail":
            return True
        if arguments.verbose: print("skipped rank " + str(i) + ". ", flush = True)
        take_screenshot(device, folder + "skipped/" + get_datetime_string() + "_skipped_" + str(i))
        return False

# Read from the power rankings list
def read_from_power_rankings(count, start, device, folder):
    if arguments.verbose: print("Reading", count, "power rankings.")
    time.sleep(1)

    i = start
    while i < (count + 1):
        if arguments.verbose: print(str(i) + "/" + str(count) + ": ", end = "", flush = True)

        # Scrape
        i = i + tap_profile(i, folder)
        read_single_power_ranking(device, folder)

        # Rest
        time.sleep(1)
        if i%10 == 0: time.sleep(5)
        if i%25 == 0: time.sleep(10)
    return

def read_top_300_stats(start, skipped, device, folder):
    if arguments.verbose: print("Reading the top 300 power rankings.")
    if skipped > 0: print("Adding", skipped, "already skipped rankings.")
    time.sleep(1)

    count = 300 + 5 + skipped
    i = start
    while i < (count + 1):
        if arguments.verbose: print(str(i) + "/" + str(count) + ": ", end = "", flush = True)

        # Scrape
        tap_modifier = tap_profile(i, folder)
        i = i + tap_modifier
        count = count + tap_modifier - 1
        grab_profile_simple(device, folder)

        # Rest
        time.sleep(1)
        if i%10 == 0: time.sleep(5)
        if i%25 == 0: time.sleep(10)
    return

def grab_profile_simple(device, folder):
    datetime = get_datetime_string()

    # Screenshot profile
    # ID, Username, Power, Kill Points
    time.sleep(2)
    take_screenshot(device, folder + datetime + "_profile")

    # Copy and save username
    device.shell("input tap 1055 455")
    time.sleep(1)
    username = get_clipboard(device)
    with open(projectFolder + "usernames.txt", 'a+', newline='', encoding='utf-8') as output_file:
        output_file.write(datetime + "\t")
        output_file.write(username)
        output_file.write("\n")
    if arguments.verbose: print(username, end = "", flush = True)

    # Exit
    if arguments.verbose: print("") # TODO: say "done" or something here if successful
    device.shell("input tap 2185 160")
    return

# ----
# Main
# ----

if __name__ == '__main__':
    # Connect to ADB device
    client = AdbClient(host="127.0.0.1", port=5037)
    device = client.device(arguments.serial)
    print("Connected to", device.serial)

    # Create project folder
    projectFolder = "output/" + arguments.project + "/"
    if not os.path.exists(projectFolder):
        os.makedirs(projectFolder + "skipped/")

    # Debug
    if arguments.debug: take_screenshot(device, projectFolder + get_datetime_string() + "_debug")

    # Scrape
    if arguments.power: read_from_power_rankings(arguments.count, arguments.start, device, projectFolder)
    if arguments.top: read_top_300_stats(arguments.start, arguments.skipped, device, projectFolder)