package com.primeton.damp.bigdata;


import lombok.extern.slf4j.Slf4j;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hive.conf.HiveConf;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.SQLException;
import java.util.Properties;
import java.util.Set;

/**
 *
 * Hive Kerberos 认证方式获得连接
 *
 *
 * <pre>
 *
 * Created by zhaopx.
 * User: zhaopx
 * Date: 2020/4/22
 * Time: 18:02
 *
 * </pre>
 *
 * @author zhaopx
 */
@Slf4j
public class Krb5HiveConnectionServiceImpl implements HiveConnectionService {


    /**
     * Hive 数据源
     */
    final Properties params;



    /**
     * 认证文件所在的基础目录
     */
    final String authBasePath;


    String hiveUrl;


    public Krb5HiveConnectionServiceImpl(Properties params) {
        this.params = params;
        this.authBasePath = params.getProperty("authBasePath");
    }


    @Override
    public boolean doAuth() {
        //KrbUser = "hadoop/cdh-node1@HADOOP.COM";
        log.info("hive 开始 kerberos 认证。");
        AuthPrincipalCreator authPrincipalCreator = AuthPrincipalCreator.useExtractorConf(authBasePath);
        Set<String> principals = authPrincipalCreator.listPrincipals();
        log.info("find existed principals: {}", principals);
        AuthPrincipal kerberosPrincipal = authPrincipalCreator.getKerberosPrincipal(params.getProperty("hiveDbUser"));

        String userKeytab = kerberosPrincipal.getUserKeytabFile().getAbsolutePath();
        String krb5File = kerberosPrincipal.getKrb5File().getAbsolutePath();
        String krbUser = kerberosPrincipal.getPrincipal();
        StringBuffer buffer = new StringBuffer(params.getProperty("hiveUrl"));
        if(!buffer.toString().contains(";principal=")) {
            buffer.append(";principal=").append(krbUser);
        }
        hiveUrl = buffer.toString();
        log.info("HIVE_URL : " + hiveUrl);

        // 分别加载 core、hdfs、hive site 文件
        Configuration conf = new Configuration();
        try {
            if (kerberosPrincipal.getCoreSite() != null) {
                conf.addResource(kerberosPrincipal.getCoreSite().toURL());
                log.info("add config: {}", kerberosPrincipal.getCoreSite().getAbsolutePath());
            }

            if (kerberosPrincipal.getHdfsSite() != null) {
                conf.addResource(kerberosPrincipal.getHdfsSite().toURL());
                log.info("add config: {}", kerberosPrincipal.getHdfsSite().getAbsolutePath());
            }

            if (kerberosPrincipal.getHiveSite() != null) {
                conf.addResource(kerberosPrincipal.getHiveSite().toURL());
                log.info("add config: {}", kerberosPrincipal.getHiveSite().getAbsolutePath());
            }
        } catch (Exception e) {
            throw new IllegalStateException(e);
        }

        // Kerberos 认证
        KerberosUtil.loginKerberos(conf, krbUser, userKeytab, krb5File);
        log.info("hive kerberos 认证通过。");
        return true;
    }

    @Override
    public Connection getConnection() throws SQLException {
        try {
            Class.forName("org.apache.hive.jdbc.HiveDriver");
        } catch (ClassNotFoundException e) {
            throw new SQLException("找不到Hive驱动：org.apache.hive.jdbc.HiveDriver.", e);
        }
        return DriverManager.getConnection(hiveUrl, "", "");
    }
}