import json,subprocess,sys
from pathlib import Path
config=json.loads(Path(__file__).with_name('access.json').read_text())
chain='CTX_METADATA_LAN'
def run(args,check=True):
 return subprocess.run(['iptables','-w','5',*args],capture_output=True,text=True,timeout=10,check=check)
rules=[['-i','lo','-j','RETURN']]+[['-s',ip+'/32','-j','RETURN'] for ip in config['readers']]+[['-j','DROP']]
jumps=[('INPUT',['-d',config['host']+'/32','-p','tcp','--dport',str(config['port']),'-j',chain]),
       ('DOCKER-USER',['-p','tcp','-m','conntrack','--ctdir','ORIGINAL','--ctorigdst',config['host'],
                       '--ctorigdstport',str(config['port']),'-j',chain])]
def verify():
 lines=run(['-S',chain]).stdout.splitlines()
 assert len([x for x in lines if x.startswith('-A ')])==len(rules)
 for rule in rules:run(['-C',chain,*rule])
 for parent,rule in jumps:run(['-C',parent,*rule])
def main(action):
 if action=='apply':
  if run(['-S',chain],False).returncode:
   run(['-N',chain])
  existing=run(['-S',chain]).stdout.splitlines()
  assert len([x for x in existing if x.startswith('-A ')])<=len(rules)
  for rule in rules:
   if run(['-C',chain,*rule],False).returncode:run(['-A',chain,*rule])
  if run(['-S','DOCKER-USER'],False).returncode:run(['-N','DOCKER-USER'])
  for parent,rule in jumps:
   if run(['-C',parent,*rule],False).returncode:run(['-I',parent,'1',*rule])
  verify()
 elif action=='check':verify()
 elif action=='remove':
  verify()
  for parent,rule in jumps:run(['-D',parent,*rule])
  for rule in rules:run(['-D',chain,*rule])
  run(['-X',chain])
 else:raise ValueError('Unknown action')
 print('METADATA_FIREWALL_'+action.upper()+'_OK')
if __name__=='__main__':main(sys.argv[1])