using BloomFilter;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Text;
using System.Text.RegularExpressions;
using System.Threading.Tasks;
using VulnCrawler;
using Newtonsoft.Json;
using Newtonsoft.Json.Linq;
using LibGit2Sharp;

namespace VulnUserCodeAnalyzer
{
    public class CVE
    {
        public string Type { get; set; }
        public int Year { get; set; }
        //public string UserName { get; set; }
        public string Code { get; set; }
        public DateTime Publish_Date { get; set; }
        public DateTime Update_Date { get; set; }
        public string Detail { get; set; }
        //public string FileName { get; set; }
        //public string FuncNameBase64 { get; set; }
        //public string Url { get; set; }
        public double Level { get; set; }
    }
    public static class CVE_JSON
    {
        /// <summary>
        /// CVE 테이블
        /// </summary>
        public static Dictionary<int, Dictionary<string, CVE>> CveDict { get; set; }
        static CVE_JSON()
        {
            CveDict = new Dictionary<int, Dictionary<string, CVE>>();
        }
        public static void AutoLoad()
        {
            var dir = new DirectoryInfo(@"c:\CVE");
            if (!dir.Exists)
            {
                Console.WriteLine("found not CVE Directory");
                return;
            }
            Console.WriteLine("Loading CVE List...");
            foreach (var json in dir.EnumerateFiles("*.json"))
            {
                var match = Regex.Match(json.Name, @"(20\d\d)");
                if (!match.Success)
                {
                    continue;
                }
                int year = int.Parse(match.Value);
  
                if (CveDict.ContainsKey(year))
                {
                    continue;
                }
                var dict = LoadCveJson(int.Parse(match.Value));
                CveDict.Add(year, dict);

                Console.WriteLine($"Finished loading CVE List Year: {year}, Count: {CveDict[year].Count}");

            }
        }

        /// <summary>
        /// CVE 정보 수집
        /// </summary>
        /// <param name="year"></param>
        /// <returns></returns>
        private static Dictionary<string, CVE> LoadCveJson(int year)
        {
            string json = File.ReadAllText($@"C:\CVE\{year}.json");
            JObject jobj = JObject.Parse(json);
            var cveDict = jobj["CVE_Items"].ToDictionary(t => t["cve"]["CVE_data_meta"]["ID"].ToString(), t =>
            {
                var vendor_data = t["cve"]["affects"]["vendor"]["vendor_data"] as JArray;
                string vendor_name = "NULL";
                if (vendor_data.Count > 0)
                {
                    vendor_name = vendor_data.First()["vendor_name"].ToString();
                }
                var description_data = t["cve"]["description"]["description_data"] as JArray;
                string description = "NULL";
                if (description_data.Count > 0)
                {
                    description = description_data.First()["value"].ToString();
                }
                double level = 0;
                var impact = t["impact"];
                if (impact.HasValues)
                {
                    level = Double.Parse(impact["baseMetricV2"]["cvssV2"]["baseScore"].ToString());
                }
                return new CVE
                {
                    Code = t["cve"]["CVE_data_meta"]["ID"].ToString(),
                    Type = vendor_name,
                    Detail = description,
                    Year = year,
                    Publish_Date = DateTime.Parse(t["publishedDate"].ToString()),
                    Update_Date = DateTime.Parse(t["lastModifiedDate"].ToString()),
                    Level = level,
                };
            });
            return cveDict;
        }
    }
    class Program
    {
        /// <summary>
        /// Clone 콜백 함수
        /// </summary>
        /// <param name="progress"></param>
        /// <returns></returns>
        public static bool TransferProgress(TransferProgress progress)
        {
            int totalBytes = progress.TotalObjects;
            int receivedBytes = progress.ReceivedObjects;
            long receivedTotal = progress.ReceivedBytes;
            double received = progress.ReceivedBytes / 1000000;
            double percent = ((double)receivedBytes / (double)totalBytes);
            Console.WriteLine($"Progress: {percent.ToString("P2")}, Remain: {receivedBytes} of {totalBytes}"); //, 받은 용량: {received.ToString()}MB");
            Console.ForegroundColor = ConsoleColor.DarkGreen;
            return true;
        }
        public static void CheckoutProcess(string path, int completedSteps, int totalSteps)
        {
            Console.WriteLine($"{completedSteps}, {totalSteps}, {path}");
        }

        public static void Clone(string path, string url)
        {
            Console.WriteLine($"Start Cloning Path : {path}");
            string clone = Repository.Clone(url, $@"{path}", new CloneOptions { OnTransferProgress = TransferProgress, OnCheckoutProgress = CheckoutProcess });
            Console.ResetColor();
            Console.WriteLine($"Finished Clone Repository: {clone}");
        }
        static void Main(string[] args)
        {
            Console.WriteLine("Start: Code Analyze Server");
            /* 연도별 CVE JSON 파일 로드 */
            CVE_JSON.AutoLoad();
            /* 크롤러 타입 */
            var crawler = new VulnC();

            /* 매칭을 위한 자료구조 Bloom Filter */
            int capacity = 50000000;
            var filter = new Filter<string>(capacity);

            /* AWS 계정 정보 파일 읽음 */
            string txt = File.ReadAllText(@"Account.xml");
            // string xml = aes.AESDecrypt128(txt, key);
            string xml = txt;
            AWS.LoadAccount(xml);
            AWS.Account account = AWS.account;
            /* AWS 정보 출력 */
            Console.WriteLine($"Endpoint: {account.Endpoint}, ID: {account.Id}");
            try
            {
                /* DB 접속 시도 */
                VulnRDS.Connect(account, "vuln");
            }
            catch (Exception e)
            {
                Console.WriteLine($"Connection Error :: {e.ToString()}");
                return;
            }
            /* AWS 연결 여부 확인 */
            if (VulnRDS.Conn.State == System.Data.ConnectionState.Open)
            {
                Console.WriteLine("Connection Success");
            }
            else
            {
                Console.WriteLine("Fail Connection");
                return;
            }
            while (true)
            {
                string userId = string.Empty;
                string repoPath = string.Empty;
                Stopwatch repoWatch = new Stopwatch();
                repoWatch.Start();
                while (true)
                {
                    var elapsedSeconds = repoWatch.Elapsed.TotalSeconds;
                    if (elapsedSeconds < 10)
                    {
                        continue;
                    }
                    Console.WriteLine("Checking User DB...");
                    var reposits = VulnRDS.SelectAllReposit();
                    foreach (var (userName, repository) in reposits)
                    {
                        if (string.IsNullOrWhiteSpace(repository))
                        {
                            continue;
                        }
                        var repoBytes = Encoding.Unicode.GetBytes(repository);
                        var repoBase64 = Convert.ToBase64String(repoBytes);
                        var repoDir = new DirectoryInfo($@"C:\Repo\{repoBase64}");
                        if (repoDir.Exists)
                        {
                            continue;
                        }
                        repoDir.Create();
                        Console.WriteLine($"Clone... Path : {repoDir.FullName}, Url : {repository}");
                        Clone(repoDir.FullName, repository);
                        repoPath = repoDir.FullName;
                        userId = userName;
                    }
                    if (!string.IsNullOrWhiteSpace(repoPath) && !string.IsNullOrWhiteSpace(userId))
                    {
                        break;
                    }
                    repoWatch.Restart();
                }
                //Console.WriteLine("엔터를 누르세요");
                //Console.ReadLine();

                /* hashDict = 사용된 사용자 함수 정보 */
                var hashDict = new Dictionary<int, HashSet<VulnAbstractCrawler.UserBlock>>();
                /* 경과 시간 체크 */
                Stopwatch stopwatch = new Stopwatch();
                stopwatch.Start();
                DirectoryInfo dirInfo = new DirectoryInfo(repoPath);

                /* 모든 .c 파일 탐색 */
                var codeFiles = dirInfo.EnumerateFiles("*.c", SearchOption.AllDirectories);
                int totalFileCount = codeFiles.Count();
                int count = 0;
                foreach (var codeFile in codeFiles)
                {
                    Console.WriteLine(codeFile.FullName);
                    using (var reader = codeFile.OpenText())
                    {
                        /* 사용자 코드를 함수별로 나눔 */
                        var dict = crawler.CrawlUserCode(reader);
                        foreach (var item in dict)
                        {
                            /* hashDict의 키와 item.key는 함수 블록의 코드 길이 */
                            if (!hashDict.ContainsKey(item.Key))
                            {
                                hashDict[item.Key] = new HashSet<VulnAbstractCrawler.UserBlock>();
                            }
                            /* item.Value는 각 코드 길이 마다의 블록 정보 
                             * Bloom Filter에 코드 블록 해쉬값 기록
                             */
                            foreach (var hash in item.Value)
                            {
                                hash.Path = codeFile.FullName;
                                hashDict[item.Key].Add(hash);
                                filter.Add(hash.Hash);
                            }
                        }
                        count++;
                        double per = ((double)count / (double)totalFileCount) * 100;
                        Console.WriteLine($"{count} / {totalFileCount} :: {per.ToString("#0.0")}%, 개체 수 : {hashDict.Count}");
                    }
                }
                var findBlocks = new Queue<VulnAbstractCrawler.UserBlock>();
                var vulnDict = new Dictionary<string, IEnumerable<VulnRDS._Vuln>>();
                foreach (var set in hashDict)
                {
                    /* 사용자 코드의 길이 마다 DB로 부터 같은 길이의 CVE 레코드 목록 가져옴 */
                    var cveList = VulnRDS.SelectVulnbyLen(set.Key).Select(v => v.Cve).Distinct();
                    foreach (var cve in cveList)
                    {
                        if (!vulnDict.ContainsKey(cve))
                        {
                            vulnDict[cve] = new HashSet<VulnRDS._Vuln>();
                            var vulnHashSet = vulnDict[cve] as HashSet<VulnRDS._Vuln>;
                            /* 같은 길이의 CVE에서 또 같은 종류의 CVE 레코드 목록 가져옴
                             * 같은 종류의 CVE 레코드들이 사용자 코드에서 모두 포함되어야 
                             * CVE를 가지고 있다고 인정하는 프로그램 정책 때문 
                             */
                            var searchedCveHashList = VulnRDS.SelectVulnbyCve(cve);

                            Console.WriteLine($"CVE:{cve}, Received Count : {searchedCveHashList.Count()}");
                            foreach (var s in searchedCveHashList)
                            {
                                vulnHashSet.Add(s);
                            }

                        }
                    }
                }
                var findCveDict = new Dictionary<string, List<VulnAbstractCrawler.UserBlock>>();
                var findCveList = new HashSet<string>();
                /* 본격적인 취약점 매칭 부분 */
                foreach (var vulnSet in vulnDict)
                {
                    Console.WriteLine($"-----cve:{vulnSet.Key}");
                    bool match = false;
                    foreach (var vuln in vulnSet.Value)
                    {
                        /* 사용자 코드 해쉬 저장해논 bloom filter에 취약점 레코드 해쉬값들이 포함되는지 확인
                         * 포함이 된다는 건 해당 취약점 레코드가 사용자 코드에도 있다는 뜻(취약점)
                         * 같은 종류의 CVE 레코드가 전부 필터에 포함된다면 취약점으로 판단한다.
                         */
                        if (filter.Contains(vuln.BlockHash))
                        {
                            if (hashDict.ContainsKey(vuln.LenFunc))
                            {
                                /* Bloom Filter는 아쉽게도 포함 여부만 알 수 있기에 
                                 * 포함되었음을 알았다면 검색해서 정보를 구한다. */
                                var userBlock = hashDict[vuln.LenFunc].FirstOrDefault(b => b.Hash == vuln.BlockHash);
                                if (userBlock == null)
                                {
                                    continue;
                                }
                                /* 해당 유저 블록을 임시 저장한다.
                                 * 밑에서 블록 정보를 DB로 전송하기 위해서다. 
                                 */
                                if (!findCveDict.ContainsKey(vuln.Cve))
                                {
                                    findCveDict[vuln.Cve] = new List<VulnAbstractCrawler.UserBlock>();
                                }
                                userBlock.Url = vuln.Url;
                                findCveDict[vuln.Cve].Add(userBlock);
                                match = true;
                            }
                        }
                        else
                        {
                            match = false;
                            break;
                        }
                    }
                    /* 취약점 레코드가 전부 있어야 CVE 찾음 인정 */
                    if (match)
                    {
                        Console.WriteLine($"Matched CVE : {vulnSet.Key}");
                        /* 찾았으면 cve값을 기록함 밑에서 찾은 cve 정보 전송하기 위해 */
                        findCveList.Add(vulnSet.Key);
                    }
                    else
                    {
                        Console.WriteLine("Not");
                    }
                }
                stopwatch.Stop();
                /* 매칭 끝 후처리 (출력, DB 전송 등) */
                var hours = stopwatch.Elapsed.Hours;
                var minutes = stopwatch.Elapsed.Minutes;
                var seconds = stopwatch.Elapsed.Seconds;
                Console.WriteLine($"Elapsed Time : {hours.ToString("00")}:{minutes.ToString("00")}:{seconds.ToString("00")}");
                Console.WriteLine($"Matched CVE Count : {findCveList.Count}");

                var yearMatch = new Regex(@"CVE-(\d{4})-(\d+)");
                foreach (var cve in findCveList)
                {
                    Console.WriteLine(cve);
                    var c = yearMatch.Match(cve);
                    int year = int.Parse(c.Groups[1].Value);
                    if (!CVE_JSON.CveDict.ContainsKey(year))
                    {
                        continue;
                    }
                    if (!CVE_JSON.CveDict[year].ContainsKey(cve))
                    {
                        continue;
                    }
                    var data = CVE_JSON.CveDict[year][cve];

                    /* 취약점 타입 분류 */
                    string type = "NORMAL";
                    if (data.Detail.IndexOf("overflow", StringComparison.CurrentCultureIgnoreCase) > 0)
                    {
                        type = "OVERFLOW";
                    }
                    else if (data.Detail.IndexOf("xss", StringComparison.CurrentCultureIgnoreCase) > 0)
                    {
                        type = "XSS";
                    }
                    else if (data.Detail.IndexOf("injection", StringComparison.CurrentCultureIgnoreCase) > 0)
                    {
                        type = "SQLINJECTION";
                    }
                    else if (data.Detail.IndexOf("dos", StringComparison.CurrentCultureIgnoreCase) > 0)
                    {
                        type = "DOS";
                    }
                    else if (data.Detail.IndexOf("Memory", StringComparison.CurrentCultureIgnoreCase) > 0)
                    {
                        type = "MEMORY";
                    }
                    else if (data.Detail.IndexOf("CSRF", StringComparison.CurrentCultureIgnoreCase) > 0)
                    {
                        type = "CSRF";
                    }
                    else if (data.Detail.IndexOf("inclusion", StringComparison.CurrentCultureIgnoreCase) > 0)
                    {
                        type = "FILEINCLUSION";
                    }
                    else if (data.Detail.IndexOf("EXCUTE", StringComparison.CurrentCultureIgnoreCase) > 0)
                    {
                        type = "EXCUTE";
                    }

                    var urlBytes = Convert.FromBase64String(findCveDict[cve].FirstOrDefault().Url);
                    string url = Encoding.Unicode.GetString(urlBytes);
                    var vulnDetail = new VulnRDS.Vuln_detail
                    {
                        CveName = data.Code,
                        Type = type,
                        Level = data.Level.ToString(),
                        Year = data.Year.ToString(),
                        CveDetail = data.Detail,
                        Publish_date = data.Publish_Date.ToString("yyyy-MM-dd"),
                        Update_date = data.Update_Date.ToString("yyyy-MM-dd"),
                        UserName = userId,
                        Url = url,
                        FileName = findCveDict[cve].FirstOrDefault().Path.Replace(repoPath, ""),
                        FuncName = findCveDict[cve].FirstOrDefault().FuncName,
                        Product = data.Type,
                    };
                    Console.WriteLine("추가 완료");
            

                    /* DB 전송 */
                    VulnRDS.InsertVulnDetail(vulnDetail);

                    Console.WriteLine($"Added CVE: {vulnDetail.CveName}, Type: {vulnDetail.Type}, CVSS: {vulnDetail.Level}");
                }
            }
        }
    }

}