記一次開源工具某模組的基礎二次開發

  • 2019 年 11 月 26 日
  • 筆記

近日看到一些關於工具的文章，很多人對於寫工具類文章都是蜻蜓點水，寫搭建，寫基礎使用，但是基礎的開源工具在生產環境上的使用其實是比較困難的，有很多需要二次開發做到與生產環境適配，今天寫這篇文章的目的就在於開啟開源工具二次開發的思路，其中一個模組的二次開發很簡單，大家基本都能看懂，就以這篇簡單易懂的程式碼改動舉例。本次開源工具二次開發以cobra程式碼審計為例。

適用範圍：程式碼初學者、有一定python基礎與協議基礎

需求明確

首先要先明確自己的需求，我們的需求是在上線前做關聯性程式碼審計，我採用的是cobra，至於為什麼要使用cobra就不多說了，這是評估期間做的事情，文章主要寫二次開發。

顯示情況

我們的生產環境使用gitlab作為程式碼庫，僅允許ssh方式下拉程式碼，而不允許http公開拉程式碼，這麼做是為了安全，但是同樣的也對cobra的使用造成了一定的麻煩。因為基礎的cobra是沒有ssh下拉功能的，需要我們自己做二次開發

程式碼改動

首先在原基礎的cobra上我們做一個拉程式碼的測試

好像很叼的樣子，沒有漏洞，但是有層script提示沒有選擇目標

後台定位問題發現如下錯誤

一般出現這個問題不是分支錯誤就是路徑錯誤，反正是找不到文件的，那直接用git clone試試原路徑是否存在吧

500，詢問了gitlab負責人才知道是根本沒有開放http下拉程式碼的功能，統一使用ssh進行下拉。

那原基礎的cobra能直接用ssh拉程式碼嗎？試試看好了

提示請輸入URL，也就是輸入格式不對，那我們換一種方式輸入

還是一樣的提示，那看看後端是什麼情況吧

一切正常。

首先對問題進行定位

命令為find ./ -name 「*.py」|xargs grep 「Please input a valid URL」

命令意思為在本級目錄級下級目錄搜索內容為Please input a valid URL的py腳本

找到了api.py，讓我們進去看看這是什麼

果然只限定了http與https，那根據咱們的需求加上ssh吧

原程式碼：

  if re.match(r'http://|https://', t):                      arg = (t, formatter, output, rule, a_sid, is_del)                      producer(task=arg)                  else:                      return {"code": 1004, "msg": "Please input a valid URL"}              result = {                  'msg': 'Add scan job successfully.',                  'sid': a_sid,                  'total_target_num': len(target),              }          else:              if re.match(r'http://|https://', target):                  arg = (target, formatter, output, rule, a_sid, is_del)                  producer(task=arg)              else:                  return {"code": 1004, "msg": "Please input a valid URL"}

改動後程式碼：

        if re.match(r'http://|https://|ssh://', t):                      arg = (t, formatter, output, rule, a_sid, is_del)                      producer(task=arg)                  else:                      return {"code": 1004, "msg": "Please input a valid URL"}              result = {                  'msg': 'Add scan job successfully.',                  'sid': a_sid,                  'total_target_num': len(target),              }          else:              if re.match(r'http://|https://|ssh://', target):                  arg = (target, formatter, output, rule, a_sid, is_del)                  producer(task=arg)              else:                  return {"code": 1004, "msg": "Please input a valid URL"}

重新試一下吧

後台提示需要輸入root密碼

我的環境是做了[email protected]的免密推送，root用戶密碼我不知道啊，就算有許可權難道把密碼寫在配置文件里又安全嗎？明文密碼泄露的事情可是發生不少，於是繼續看程式碼找問題

find ./ -name "*.py" |xargs grep "git clone"

看他的推送方式是什麼

配置文件為./cobra/pickup.py

如果配置文件中沒有輸入用戶名或密碼，便是公開鏈接，直接clone，如果有用戶名密碼，則分割填入用戶名和密碼進行加密clone，如果這放在http協議中這個邏輯完全沒有問題，但是放在ssh下拉程式碼里，就會存在很大的邏輯問題了

邏輯問題在於ssh的免密鑰登錄不需要密碼，如果單純以用戶名密碼作為判斷依據那免密鑰的作用就為零了，為了適配免密鑰，實際上在程式碼中指定免密鑰的用戶即可，同時為保證其他基礎功能不遭受破壞，增加的功能應使用判斷前綴的方式進行程式碼改寫，具體如下：

原程式碼：

         if self.repo_username is None or self.repo_password is None:              # public repo              clone_address = self.repo_address          else:              # private repo              clone_address = self.repo_address.split('://')[0] + '://' + quote(self.repo_username) + ':' +                               quote(self.repo_password) + '@' + self.repo_address.split('://')[1]          # clone repo with username and password            # "http[s]://username:[email protected]/username/reponame"          # !!! if add password in the url, .git/config will log your url with password          cmd = 'git clone ' + clone_address + ' "' + self.repo_directory + '" -b ' + self.repo_branch

改後程式碼：

     if self.repo_username is None or self.repo_password is None:              # public repo              if (self.repo_address.split('://')[0] == 'ssh'):                  clone_address = 'ssh://' + 'git@' +                   self.repo_address.split('://')[1]              else:              clone_address = self.repo_address          else:              # private repo              if (self.repo_address.split('://')[0] == 'ssh'):                  clone_address = 'ssh://' + 'git@' +                   self.repo_address.split('://')[1]              else:              clone_address = self.repo_address.split('://')[0] + '://' + quote(self.repo_username) + ':' +                               quote(self.repo_password) + '@' + self.repo_address.split('://')[1]          # clone repo with username and password          # "http[s]://username:[email protected]/username/reponame"          # !!! if add password in the url, .git/config will log your url with password          cmd = 'git clone ' + clone_address + ' "' + self.repo_directory + '" -b ' + self.repo_branch

增加前綴判斷，如果前綴為ssh，則統一使用通用用戶git進行ssh登錄程式碼拉取

進行測試

如此完成了一次很簡單的模組二次開發。

二次開發重要是需求與實現，將自己的需求完整的實現出來即可，不一定是框架意義上的開發才算二次開發，一定要明白自己需要的是什麼，程式碼邏輯是什麼，掌握好程式碼邏輯，二次開發也不算是很難的事情。

*本文原創作者：煜陽yuyang，本文屬於FreeBuf原創獎勵計劃，未經許可禁止轉載

VirMach 便宜 VPS

QNews

QNews