valconomy/controller/valconomy.py

import base64
from dataclasses import dataclass
from enum import Enum
import errno
import json
import logging
from pprint import pprint
import os
import queue
import struct
import time
from typing import Callable

import hid
import requests
import urllib3

# Local HTTPS connection, we need to not verify the cert
urllib3.disable_warnings()
# Disable so we can debug without a log for every request
logging.getLogger('urllib3.connectionpool').disabled = True

log = logging.getLogger('valconomy')

@dataclass
class ValorantPlayerInfo:
  is_idle: bool = False

  is_party_owner: bool = True
  max_party_size: int = 5
  party_size: int = 1
  party_state: str = None

  queue_type: str = None
  game_state: str = None

  map: str = None
  team: str = None
  score: int = 0
  enemy_score: int = 0

@dataclass
class RiotPlayerInfo:
  uuid: str
  name: str
  tag: str
  state: str

  valorant: ValorantPlayerInfo = None

  @classmethod
  def dummy(cls, **kwargs):
    return cls('00000000-0000-0000-0000-000000000000', 'Player', 'gamer', 'dnd', **kwargs)

  def full_name(self) -> str:
    return f'{self.name}#{self.tag}'

class ValorantLocalClient:
  poll_interval = 0.5

  def __init__(self, callback: Callable[[RiotPlayerInfo, bool], None]):
    self.callback = callback

    self.lockfile_path = os.path.join(os.environ['LOCALAPPDATA'], r'Riot Games\Riot Client\Config\lockfile')
    self.port, self.password = None, None
    self.running = True
    self.players = {}

  def _load_credentials(self):
    if not os.path.exists(self.lockfile_path):
      return False

    try:
      with open(self.lockfile_path) as f:
        for line in f:
          if line.startswith('Riot Client'):
            client, pid, self.port, self.password, proto = line.rstrip().split(':')
            assert proto == 'https'
            return True
    except FileNotFoundError:
      return False

    return False

  def _do_poll(self):
    try:
      resp = requests.get(
        f'https://127.0.0.1:{self.port}/chat/v4/presences',
        auth=('riot', self.password), verify=False)
      resp.raise_for_status()
      data = resp.json()
    except (requests.HTTPError, requests.ConnectionError) as ex:
      # Most likely one of these means the Riot Client shut down
      log.error(f'Failed to make request to Riot Client: {ex}')
      return False

    next_players = {}
    for p in data['presences']:
      if p['product'] != 'valorant':
        continue

      v = json.loads(base64.b64decode(p.pop('private'))) # Valorant-specfic
      # pprint(p)
      # pprint(v)

      val_info = ValorantPlayerInfo(
        is_idle=v['isIdle'], is_party_owner=v['isPartyOwner'],
        max_party_size=v['maxPartySize'], party_size=v['partySize'],
        party_state=v['partyState'],
        queue_type=v['queueId'] or None, game_state=v['sessionLoopState'],
        map=v['matchMap'] or None,
        score=v['partyOwnerMatchScoreAllyTeam'], enemy_score=v['partyOwnerMatchScoreEnemyTeam'])
      match v['partyOwnerMatchCurrentTeam']:
        case 'Red':
          val_info.team = 'attackers'
        case 'Blue':
          val_info.team = 'defenders'

      info = RiotPlayerInfo(p['puuid'], p['game_name'], p['game_tag'], p['state'], valorant=val_info)

      last = self.players.get(info.uuid)
      if info != last:
        self.callback(info, False)
      next_players[info.uuid] = info

    for uuid in self.players.keys() - next_players.keys():
      self.callback(self.players[uuid], True)

    self.players = next_players
    return True

  def run(self):
    while self.running:
      if self.password is None:
        if not self._load_credentials():
          time.sleep(self.poll_interval)
          continue
        log.info('Detected Riot Client credentials, starting polling')

      if not self._do_poll():
        for p in self.players.values():
          self.callback(p, True)
        self.players = {}
        self.port, self.password = None, None

      time.sleep(self.poll_interval)

  def stop(self):
    self.running = False

class EconomyDecision(Enum):
  BUY = 0
  SAVE = 1
  BONUS = 2
  MATCH_TEAM = 3

class ValconomyHandler:
  # Valorant isn't open
  def none(self):
    log.info('Val Time soon?')

  # Welcome the user (back) to the game!
  def menu(self, info: RiotPlayerInfo, was_idle: bool=False):
    if was_idle:
      log.info(f"Welcome back, {info.name}!")
    else:
      log.info(f"It's Val Time, {info.name}!")

  # The user is idle in the menu...
  def idle(self, info: RiotPlayerInfo):
    log.info(f'Come back soon, {info.name}...')

  # Just entered queue
  def queue_start(self, info: RiotPlayerInfo):
    if info.valorant.is_party_owner and info.valorant.queue_type == 'unrated':
      log.info('Uhhhh should that be comp, MoonStar?')
    else:
      log.info(f'Hope you find a game quickly, {info.name}!')

  def match_found(self, info: RiotPlayerInfo):
    if info.valorant.queue_type == 'premier-seasonmatch':
      log.info('Do the Cosmonauts proud!')
    else:
      log.info("I hope it's not Split...")

  # Loaded into the agent select
  def pregame(self, info: RiotPlayerInfo):
    if info.valorant.map == '/Game/Maps/Bonsai/Bonsai':
      log.info('Ewwwww, Split....')
    else:
      log.info('Pick a good agent!')

  # Game where we're not providing economy help
  def game_generic(self, info: RiotPlayerInfo):
    match info.valorant.queue_type:
      case 'hurm':
        log.info('Have a good TDM!')
      case 'deathmatch':
        log.info('Have a good deathmatch!')
      case 'ggteam':
        log.info('Have a good Escalation!')
      case 'spikerush':
        log.info('Have a good Spike Rush!')
      case _:
        if info.valorant.max_party_size == 12:
          log.info('Have a good custom!')
        else:
          log.info('Have a good game!')

  # Loaded into the game
  def game_start(self, info: RiotPlayerInfo):
    log.info('OK best of luck!')

  # Round started
  def round_start(self, info: RiotPlayerInfo, won: bool, economy: EconomyDecision):
    msg = f'Start round {info.valorant.score}-{info.valorant.enemy_score}'
    if won is not None:
      msg += f' ({"won" if won else "lost"})'
    log.info(msg)

    match economy:
      case EconomyDecision.BUY:
        log.info('You should buy!')
      case EconomyDecision.SAVE:
        log.info('Time to save...')
      case EconomyDecision.BONUS:
        log.info('Keep your gun from last round!')
      case EconomyDecision.MATCH_TEAM:
        log.info('Follow the team economy!')

  def game_over(self, info: RiotPlayerInfo, won: bool):
    if won:
      log.info('Well played!')
    else:
      log.info('Hard luck...')

class HIDValconomyHandler(ValconomyHandler):
  vid = 0x6969
  pid = 0x0004

  def __init__(self):
    self.dev = None
    self.running = True
    self.queue = queue.Queue(128)

  def close(self):
    if self.dev is not None:
      self.dev.close()

  def _dev_ready(self):
    try:
      if self.dev is None:
        dev = hid.device()
        dev.open(vendor_id=self.vid, product_id=self.pid)
        self.dev = dev
        log.info(f'USB device opened')

      # 2 bytes: report ID and returned value
      # We get back the same report ID and value
      # Set to report size + 1 to satisfy W*ndoze
      data = self.dev.get_input_report(0, 65)
      assert len(data) == 2
      return data[1] == 1
    except OSError as ex:
      if self.dev is not None:
        log.warning(f'USB device lost')
        self.dev = None
      return False

  def _do(self, cmd, *vals, fmt=None):
    if fmt is None:
      fmt = ''

    fmt = '<BB' + fmt
    # Prepend report ID 0
    data = struct.pack(fmt, *(0, cmd) + vals)
    self.dev.write(data)

  def _enq(self, cmd, *vals, fmt=None):
    self.queue.put((cmd, vals, fmt))

  def service(self):
    while not self.queue.empty():
      cmd, vals, fmt = self.queue.get()

      while not self._dev_ready():
        if not self.running:
          break
        time.sleep(0.1)

      try:
        self._do(cmd, *vals, fmt=fmt)
      except OSError as ex:
        log.warning(f'USB device lost, state dequeuing stalled')

  def run(self):
    while self.running:
      while self.queue.empty():
        time.sleep(0.5)

      self.service()

  def none(self):
    self._enq(0)

  def menu(self, info: RiotPlayerInfo, was_idle: bool=False):
    self._enq(1, 1 if was_idle else 0, fmt='B')

  def idle(self, info: RiotPlayerInfo):
    self._enq(2)

  def queue_start(self, info: RiotPlayerInfo):
    ms_not_comp = info.valorant.is_party_owner and info.valorant.queue_type == 'unrated'
    self._enq(3, 1 if ms_not_comp else 0, fmt='B')

class GameState(Enum):
  NONE = 0
  MENU = 1
  IDLE = 2
  IN_QUEUE = 3
  GAME_FOUND = 4
  PRE_GAME = 5
  GAME_START = 6
  IN_GAME = 7
  GAME_OVER = 8
  IN_GAME_GENERIC = 9

class ValconomyStateMachine:
  def __init__(self, player_uuid: str, handler: ValconomyHandler):
    self.uuid = player_uuid
    self.handler = handler

    self.last = None
    self.score = -1
    self.score_enemy = -1
    self.round_history = []

  def handle_state(self, s: GameState, p: RiotPlayerInfo):
    log.info(f'{self.last} -> {s}')

    if s == GameState.IN_GAME:
      if p.valorant.score == self.score and p.valorant.enemy_score == self.score_enemy:
        # Same score line
        return

      won_last = None
      if self.score != -1:
        won_last = True if p.valorant.score > self.score else False
        self.round_history.append(won_last)

      if p.valorant.queue_type == 'swiftplay':
        if p.valorant.score == 5 or p.valorant.enemy_score == 5:
          # Game is over
          return

        eco = EconomyDecision.BUY
      else:
        if p.valorant.score > 12 or p.valorant.enemy_score > 12:
          if p.valorant.queue_type == 'unrated' or abs(p.valorant.score - p.valorant.enemy_score) >= 2:
            # Match is over
            return

        eco = EconomyDecision.MATCH_TEAM
        rounds_played = p.valorant.score + p.valorant.enemy_score
        if p.valorant.enemy_score == 12:
          # Enemy about to win!
          eco = EconomyDecision.BUY
        elif rounds_played in (0, 12):
          # First round (of half)
          eco = EconomyDecision.BUY
        elif rounds_played in (1, 13):
          # Second round (of half)
          eco = EconomyDecision.BUY if won_last else EconomyDecision.SAVE
        elif rounds_played in (2, 14):
          # Third round (of half)
          match self.round_history[-2:]:
            case [True, True]:
              eco = EconomyDecision.BONUS
            case [True, False]:
              eco = EconomyDecision.SAVE
            case [False, _]:
              eco = EconomyDecision.BUY
        elif rounds_played >= 24:
          # Sudden death or overtime (buy either way)
          eco = EconomyDecision.BUY

      self.handler.round_start(p, won_last, eco)
      self.score = p.valorant.score
      self.score_enemy = p.valorant.enemy_score
      self.last = GameState.IN_GAME
      return

    if s == self.last:
      return

    if s != GameState.GAME_OVER:
      self.score = self.score_enemy = -1
      self.round_history = []
    match s:
      case GameState.IN_GAME_GENERIC:
        self.handler.game_generic(p)
      case GameState.GAME_START:
        self.handler.game_start(p)
      case GameState.PRE_GAME:
        self.handler.pregame(p)
      case GameState.GAME_FOUND:
        self.handler.match_found(p)
      case GameState.IN_QUEUE:
        self.handler.queue_start(p)

      case GameState.IDLE:
        self.handler.idle(p)
      case GameState.MENU:
        self.handler.menu(p, self.last == GameState.IDLE)

      case GameState.NONE:
        self.handler.none()

      case GameState.GAME_OVER:
        self.handler.game_over(p, self.score > self.score_enemy)

    self.last = s

  def handle_presence(self, p: RiotPlayerInfo, gone: bool):
    log.debug(f'{repr(p)}, gone: {gone}')
    if p.uuid != self.uuid:
      return

    if gone:
      self.handle_state(GameState.NONE, p)
      return

    if p.valorant.party_state == 'MATCHMADE_GAME_STARTING':
      self.handle_state(GameState.GAME_FOUND, p)
      return
    if p.valorant.party_state == 'MATCHMAKING':
      self.handle_state(GameState.IN_QUEUE, p)
      return

    if p.valorant.queue_type in ('unrated', 'competitive', 'premier-seasonmatch', 'swiftplay'):
      if p.valorant.game_state == 'PREGAME':
        self.handle_state(GameState.PRE_GAME, p)
        return
      if p.valorant.game_state == 'INGAME':
        if self.score == -1 and p.valorant.score == p.valorant.enemy_score == 0:
          self.handle_state(GameState.GAME_START, p)

        if (self.score > 0 or self.score_enemy > 0) and p.valorant.score == p.valorant.enemy_score == 0:
          self.handle_state(GameState.GAME_OVER, p)
        else:
          self.handle_state(GameState.IN_GAME, p)
        return

    if p.valorant.game_state in ('PREGAME', 'INGAME'):
      self.handle_state(GameState.IN_GAME_GENERIC, p)
      return

    if p.valorant.game_state == 'MENUS':
      if p.valorant.is_idle:
        self.handle_state(GameState.IDLE, p)
        return

      self.handle_state(GameState.MENU, p)