#!/usr/bin/python3
import rados
import sys
import json
import argparse

parser = argparse.ArgumentParser(prog="ct-get-degraded-rbd", description="Script to show degraded RBD")
parser.add_argument("-p", "--pools", action="store",default="nova", type=str, help="Pools to analyse separate by comma (default: nova)")
parser.add_argument("-m", "--max_pgs", action="store",default=1024, type=int, help="Max PGs to analyse by pool (default: 1024)")
parser.add_argument("-s", "--state", action="store", default="inactive", type=str, help="State of PGs to analyse (default: inactive)")
parser.add_argument("-j", "--json", action="store_true", help="Output in json format")
args = parser.parse_args()

pools_str = args.pools.split(',')

metadatas = set()
datas = set()
pools = {}
out_json = {}
cluster = rados.Rados(conffile='/etc/ceph/ceph.conf')
cluster.connect()
cmd = json.dumps({'prefix': 'osd lspools', 'format': 'json'})
ret,stdout,_ = cluster.mon_command(cmd,b'')
pools_dict = json.loads(stdout)
for pool in pools_dict:
  if pool['poolname'] in pools_str:
    pools[pool['poolname']] = pool['poolnum']

state = args.state # down,laggy,degraded,repair,recovery_unfound,backfill_unfound,incomplete,stale,inconsistent

cmd = json.dumps({'prefix': "pg dump_stuck", 'states': ['clean'], 'format': 'json', '': 2})
ret,stdout,_ = cluster.mon_command(cmd,b'')
pgs = json.loads(stdout)
if 'pg_stats' not in pgs:
  quit()
pges = [ x['pgid'] for x in pgs['pg_stats']]

for pool,poolid in pools.items():
  cpt_pgs = 0
  rbd_datas = set()
  ioctx = cluster.open_ioctx(pool)
  for pg in pges:
    if pg.startswith(f"{poolid}."):
      cpt_pgs += 1
      if cpt_pgs > args.max_pgs:
        print(f"Max PGs to analyse on {pool} (>{args.max_pgs})")
        out_json['Max PGs to analyse'] = f"on {pool} (>{args.max_pgs})"
        break
      ioctx.set_locator_key(pg)
      objects = ioctx.list_objects()
      for object in objects:
        name = object.key.split('.')
        if name[0] == 'rbd_id':
          metadatas.add(name[1])
        elif name[0] == 'rbd_data':
          rbd_datas.add(name[1])
  ioctx.set_locator_key('')
  for rbd_data in rbd_datas:
    key = f"id_{rbd_data}"
    with rados.ReadOpCtx(ioctx) as read_op:
      try:
        iter, ret = ioctx.get_omap_vals_by_keys(read_op,(key,))
        ioctx.operate_read_op(read_op, "rbd_directory")
        name = list(iter)[0][1][4:].decode('utf8')
        datas.add(name)
      except:
        pass
  ioctx.close()
cluster.shutdown()

if args.json:
  out_json['affected metadatas'] = list(metadatas)
  out_json['affected datas'] = list(datas)
  print(out_json)
else:
  print("Affected metadatas:")
  for md in metadatas:
    print(f"\t{md}")
  print("Affected data:")
  for d in datas:
    print(f"\t{d}")
