#!/usr/bin/env python
# coding:utf8

import sys
import os
import shutil
import time
import re
import json

if os.path.exists(‘tab.py‘):
    import tab

class handleHaproxyConf(object):
    "search, add, update and delete functions."
    def __init__(self, lines):
        self.lines = lines

    def check(self, b):
        lines = self.lines.copy()
        for i in lines:
            if re.match(r‘^backend %s$‘ % b, i): 
                return True
        return False

    def check_dict(self, d):
        backend = d.get(‘backend‘)
        record = d.get(‘record‘)
        if backend and record:
            if record.get(‘server‘) and record.get(‘weight‘) and record.get(‘maxconn‘):
                return True
            else:
                return False
        else:
            return False

    def search(self, b):
        sList = []
        lines = self.lines.copy()
        for i in lines:
            if re.match(r‘^backend %s$‘ % b, i):
                for k in lines[(lines.index(i)+1):]:
                    if not re.search(r‘^\s+server‘, k):
                        break
                    sList.append(k)
        return sList

    def add(self, n):
        if not self.check_dict(n):
            return False
        if self.check(n.get(‘backend‘)):
            return self.update(n)
        else:
            a = []
            lines = self.lines.copy()
            lines.append(‘backend ‘ + n.get(‘backend‘) + ‘\n‘)
            for i in n.get(‘record‘).items():
                for k in i:
                    a.append(str(k))
            a.insert(1, a[1])
            lines.append(‘\t‘ + ‘ ‘.join(a) + ‘\n‘)
            return lines

    def update(self, n):
        if not self.check_dict(n):
            return False
        lines = self.lines.copy()
        for i in self.lines.copy():
            if re.match(r‘^backend %s$‘ % n.get(‘backend‘), i): 
                for k in lines[(lines.index(i)+1):]:
                    if re.search(r‘^\s+server‘, k): 
                        klist = k.strip(‘ |\t|\n‘).split(‘ ‘)
                        if klist[1] == n.get(‘record‘).get(‘server‘):
                            print(‘already exists.‘)
                            return False
                        klist[1:3] = [str(n.get(‘record‘).get(‘server‘))] * 2
                        klist[4] = str(n.get(‘record‘).get(‘weight‘))
                        klist[6] = str(n.get(‘record‘).get(‘maxconn‘))
                        lines[lines.index(k)] = ‘\t‘ + ‘ ‘.join(klist) + ‘\n‘
                        return lines
        return False

    def delete(self, n):
        d = []
        lines = oldlines = self.lines.copy()
        if self.check(n):
            for i in oldlines:
                if re.match(r‘^backend %s$‘ % n, i):    
                    s = oldlines.index(i)
                    lines.pop(s)
                    for k in oldlines[s:]:
                        if re.search(r‘^\s+server‘, k):
                            lines.pop(s)
                        else:
                            break
            return lines
        else:
            return False

def main():
    
    def writeConf(lines, hfile):
        shutil.copy2(hfile, ‘%s.%s‘ % (hfile, time.strftime("%Y%m%d_%H%M", time.localtime())))
        with open(hfile, ‘w‘) as f:
            f.write(‘‘.join(lines))
        print("success.")

    example1 = ‘www.centos.org‘
    example2 = {‘backend‘: ‘www.centos.org‘, ‘record‘: {‘server‘: ‘100.1.1.1‘, ‘weight‘: 20, ‘maxconn‘: 3000}}
    display = "1) Search\n2) Add\n3) Update\n4) Delete\nq) Exit"
    print("--Handle haproxy.conf--\n%s" % display)

    while True:
        try:
            haproxyConf = ‘haproxy.conf‘
            with open(haproxyConf, ‘r‘) as f:
                lines = f.readlines()
            
            hhc = handleHaproxyConf(lines)
        
            choice = input(">> ").strip()
            if choice == ‘1‘: # search
                print("example: %s\n" % example1)
                s = input("Search backend: ").strip()
                s_result = hhc.search(s)
                if s_result:
                    for i in s_result:
                        print(i.strip(‘\n‘))
                else:
                    print("not found...")
            elif choice == ‘2‘: # add
                print("example: %s\n" % example2)
                a = eval(input("Add backend (dict): "))
                a_result = hhc.add(a)
                #print(‘‘.join(a_result))
                if a_result:
                    writeConf(a_result, haproxyConf)
                else:
                    print("format error..")
                    continue
            elif choice == ‘3‘: # update
                print("example: %s\n" % example2)
                u = eval(input("Update backend (dict): "))
                u_result = hhc.update(u)
                if u_result:
                    writeConf(u_result, haproxyConf)
                else:
                    print("not found...")
            elif choice == ‘4‘: # delete
                print("example: %s\n" % example1)
                d = input("Delete backend: ")
                d_result = hhc.delete(d)
                if d_result:
                    writeConf(d_result, haproxyConf)
                else:
                    print("not found...")
            elif choice == ‘q‘: # exit
                print("exit...")
                break
            continue
        except:
            print("error...\n")   
            print(display)
            continue

if __name__ == ‘__main__‘:
    main()