#! /usr/bin/python3
import rados
import sys
import json
import argparse
from subprocess import run

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=32, type=int, help="Max PGs to analyse by pool (default: 32)")
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()
c = run(["ceph","osd","lspools","-f","json"], capture_output=True)
pools_dict = json.loads(c.stdout)
for pool in pools_dict:
  if pool['poolname'] in pools_str:
    pools.append(pool)

state = args.state # down,laggy,degraded,repair,recovery_unfound,backfill_unfound,incomplete,stale,inconsistent
for pool in pools:
  cpt_pgs = 0
  rbd_datas = set()
  c = run(["ceph","pg","dump_stuck",str(pool['poolnum']),"-f","json",state ], capture_output=True)
  try:
    pgs = json.loads(c.stdout)
    if 'pg_stats' in pgs:
      for pg in pgs['pg_stats']:
        cpt_pgs += 1
        if cpt_pgs > args.max_pgs:
          print(f"Max PGs to analyse on {pool['poolname']} (>{args.max_pgs})")
          out_json['Max PGs to analyse'] = f"on {pool['poolname']} (>{args.max_pgs})"
          break
        c = run(["rados","-p",pool['poolname'],'--pgid',pg['pgid'],"-f","json",'ls'], capture_output=True)
        objects = json.loads(c.stdout)
        for object in objects:
          name = object['name'].split('.')
          if name[0] == 'rbd_id':
            metadatas.add(name[1])
          elif name[0] == 'rbd_data':
            rbd_datas.add(name[1])
    for rbd_data in rbd_datas:
      ioctx = cluster.open_ioctx(pool['poolname'])
      key = f"id_{rbd_data}"
      with rados.ReadOpCtx(ioctx) as read_op:
        iter, ret = ioctx.get_omap_vals_by_keys(read_op,(key,))
        ioctx.operate_read_op(read_op, "rbd_directory")
        datas.add(list(iter)[0][1][4:].decode('utf8'))
      ioctx.close()
  except:
    pass
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}")
