【第一只爬虫】统计Github项目中的目录/文件个数

被迫转岗，开始接触QA的工作。本系列文章主要记录一些工作日常中踩到的坑和接触到的新知识，把握好这段短暂的时光吧～

今天是值得给自己鼓鼓掌的一天～好不容易写出了人生中的一个爬虫！

以上写在2020.1.21

需求说明

统计Github项目中的规则数量

1. red-team
   1. windows系统：xxx条
   2. linux系统：xxx条
2. purple-team
   1. windows系统：xxx条
   2. linux系统：xxx条

思路整理

首先分析一下Github两个项目的目录结构

1. red-team项目是以“规则”作为分类的依据，目录名为“T”+4位数字。
2. 每个规则目录下至少包含一个同名的markdown文件和一个同名的yaml文件。
3. yaml文件中存放了多个攻击手段，根据需求方要求，每一种攻击手段我们都视为一条规则。
4. 每种攻击手段都标记出了支持的操作系统有哪些。

1. purple-team项目是以“支持平台”作为分类的依据，目录名分别为“windows”和“linux”。
2. windows 目录下存放了文件名为“t”+4位数字+“.rb”的规则文件和其他的一些说明文件（无需统计）。
3. linux 目录下存放了文件名为“t”+4位数字+“.rb”的规则文件。

明确统计思路

1. red-team可以通过遍历所有文件名是以“T”+4位数字开头的yaml文件，统计windows和linux字符串分别累计出现了多少次
2. purple-team可以分别统计windows和linux目录下，文件名是以“t”+4位数字开头的文件数量。

既是Github上的开源项目，我们可以将项目克隆到本地（方案一），也通过爬虫直接统计（方案二）。接下来将从方案二展开叙述。

实现过程

导入第三方库

import requests
import re
from bs4 import BeautifulSoup
from tqdm import tqdm

requests 模块可以模拟浏览器的请求
re 模块通过正则表达式匹配，对字符串进行处理
bs4 库的 BeautifulSoup 模块可以快速提取网页上的内容，熟悉HTML/CSS的话理解起来更快。
tqdm 模块可以在终端执行时，展示进度，对于需要长时间运行的脚本有非常好的用户体验。

源码

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import requests
import re
from bs4 import BeautifulSoup
from tqdm import tqdm

def count_purple_team_rules(keywords):
    url = 'https://github.com/praetorian-code/purple-team-attack-automation/tree/master/modules/post/' + keywords +'/purple'
    headers = {"User-Agent": "Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/68.0.3440.106 Safari/537.36"}
    res1 = requests.get(url, headers=headers)
    soup = BeautifulSoup(res1.text,'lxml')
    # 查找所有包含title属性的节点
    node_list = soup.find_all(title=re.compile('t[0-9]{4}'))
    # 计算长度
    return len(node_list)

def get_red_team_rule_ids():
    url = 'https://github.com/redcanaryco/atomic-red-team/blob/master/atomics'
    headers = {"User-Agent": "Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/68.0.3440.106 Safari/537.36"}
    res1 = requests.get(url, headers=headers)
    soup = BeautifulSoup(res1.text,'lxml')
    id_list = []
    for item in soup.find_all(title=re.compile('^T[0-9]{4}$')):
        context = item.string
        id_list.append(context)
    return id_list

def count_key(id_list):
    win_count = 0
    linux_count = 0
    for id in tqdm(id_list,desc='正在爬取页面数据'):
        url = 'https://github.com/redcanaryco/atomic-red-team/blob/master/atomics/'+ id + '/' + id + '.yaml'
        headers = {"User-Agent": "Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/68.0.3440.106 Safari/537.36"}
        res1 = requests.get(url, headers=headers)
        soup = BeautifulSoup(res1.text,'lxml')
        for item in soup.find_all(class_='pl-s'):
            context = item.string
            if context == 'windows':
                win_count = win_count + 1
            if context == 'linux':
                linux_count = linux_count + 1
    return win_count, linux_count

if __name__ == "__main__":
    red_team_id_list = get_red_team_rule_ids()
    print(red_team_id_list)
    count_turple = count_key(id_list)
    print(count_turple[0])
    print(count_turple[1])

总结

功能可以实现，但仍有很多值得优化的地方：

1. 模块层次还不够明晰。
2. 有一些重复适用的语句可以单独写为一种方法。
3. 执行时间比较长，5min才可以输出结果。

希望在后续接触更多的脚本之后，积累了更多的认识，能把以上的问题解决了。