import mysql.connector
import sys


class DBhelper(object):
    """
    可能的问题：
    每次进行连接会不会造成时延呢？
    这里没有进行关闭喔
    一条数据一条数据的更新效率可能不高
    """
    def __init__(self, host='localhost', port=3306, user='root',
                 database="test"):  # 构造函数
        self.tablename = 'flow_feature'
        try:
            self.conn = mysql.connector.connect(user='root', host='localhost',password='password',
                                                database='packetin_flow',buffered=True)
          
            self.cursor = self.conn.cursor()
            print('Connected', self.conn)
            self.cursor.execute("SHOW TABLES")
        except Exception as e:
            print(e)
            print('gg!')

    def create_table(self):
        print('Init table', self.tablename)
        self.cursor.execute("DROP TABLE IF exists %s" % self.tablename)
        self.cursor.execute("""CREATE TABLE %s (
          `id` INT unsigned NOT NULL AUTO_INCREMENT,
          `ipSet` varchar(128) NOT NULL,
          `dst_port` INT NOT NULL,
          `src_port` INT NOT NULL,
          `flags` INT NOT NULL,
          `window_size` INT NOT NULL,
          `reason` INT NOT NULL,
          `in_port` INT NOT NULL,
          `min_len` INT NOT NULL,
          `max_len` INT NOT NULL,
          `means_len` FLOAT NOT NULL,
          `pcapnum` INT NOT NULL,
          `tagY` INT NOT NULL,
          PRIMARY KEY (`id`)
          ) ENGINE=InnoDB DEFAULT CHARSET=utf8""" % self.tablename)
        return 1

    def delete_table(self):
        print('Delete table', self.tablename)
        self.cursor.execute("DROP TABLE IF EXISTS %s" % self.tablename)
        return 1

    def execute(self, sql):
        self.cursor.execute(sql)
        rowcount = self.cursor.rowcount
        return rowcount

    def get_ipset_value(self, ipSet):
        """
        查询函数
        :param ipSet: 传入的是一个tuple类型的ipset
        :return: 返回对应的dict，如果有多条记录只返回一条，如果查询失败返回0
        """
        ipSet_str = ' '.join(ipSet)
        sql = """select * from %s where ipSet='%s' limit 1""" % (self.tablename, ipSet_str)
        try:
            # 执行SQL语句
            self.cursor.execute(sql)
            # 使用 fetchone() 方法获取单条数据.
            data = self.cursor.fetchone()
        except:
            # 发生错误时回滚
            print('select error, now rollback...')
            self.conn.rollback()
            return 0
        this_feature = list(data)
        this_feature.pop(0)         # 把id和ipSet删了
        this_feature.pop(0)
        return {ipSet: this_feature}

    def judgeIPset(self, ipSet):
        # select 1 from tablename where col = 'col' limit 1
        # select isnull((select top(1) 1 from %s where `ipSet`=%s), 0)
        ipSet_str = ' '.join(ipSet)
        sql = """
            select * from %s where ipSet = '%s' limit 1
            """ % (self.tablename, ipSet_str)
        out = self.cursor.execute(sql)
        out = self.cursor.rowcount              # 不知为何用execute的out不行
        if out:
            print('ipSet in database.')
        else:
            print('Not in database, now insert...')
        return out

    def insert(self, ipSet, feature_list):
        """
        插入数据进mysql
        :param ipSet:
        :param feature_list: 传入的数据是处理好pcapnum和tagY的
        :return: 0失败 1成功
        """
        ipSet_str = ' '.join(ipSet)
        sql = """
            INSERT INTO {0} 
            (ipSet, dst_port, src_port, flags, window_size, 
            reason, in_port, min_len, max_len, means_len, pcapnum, tagY)
            VALUES
            (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        """.format(self.tablename)
        values = (ipSet_str,) + tuple(feature_list)
        try:
            self.cursor.execute(sql, values)
            self.conn.commit()
            print(self.cursor.rowcount, 'Insert successful', self.cursor.lastrowid)
            # self.cursor.close()
            # self.conn.close()
            return self.cursor.lastrowid
        except:
            # Rollback in case there is any error
            print('Insert error, now rollback...')
            self.conn.rollback()
            return 0

    def update(self, feature_list, ipSet):
        """
        数据库中本身存在，更新预测结果
        :param feature_list: 传入一个list，已经更新好的喔
        :param ipSet
        :return:
        """
        ipSet_str = ' '.join(ipSet)
        sql = """
            UPDATE {0} SET
            `ipSet` = %s, `dst_port` = %s, `src_port` = %s, `flags` = %s, `window_size` = %s,
            `reason` = %s, `in_port` = %s, `min_len` = %s, `max_len` = %s, `means_len` = %s,
            `pcapnum` = %s, `tagY` = %s
            where ipSet='{1}'
            """.format(self.tablename, ipSet_str)
        values = (ipSet_str,) + tuple(feature_list)
        try:
            self.cursor.execute(sql, values)
            self.conn.commit()
            return 1
        except:
            # Rollback in case there is any error
            print('更新失败，要哭了')
            self.conn.rollback()
            return 0


if __name__ == '__main__':
    """
    需要先打开mysql服务，并确保有database名为packetin_flow
    在执行packetinpredict的时候，需要先建表
    mydb.create_table()
    """
    mydb = DBhelper()
    mydb.create_table()
    out = mydb.execute("""
    select * from flow_feature where ipSet = '188.42.254.65 10.0.2.107' limit 1
    """)
    print('rowcount: ', mydb.cursor.rowcount)
    # mydb.delete_table()