From 6d445bcbd286209582eaa3f8250a96ebf09afa93 Mon Sep 17 00:00:00 2001
From: Everardo Padilla <everardo.padilla@wolt.com>
Date: Mon, 29 Nov 2021 18:04:57 +0200
Subject: [PATCH] Add support for pipenv/pipfile

---
 README.md           | 12 +++++++++--
 main.go             |  6 +++---
 pip.go => python.go | 49 ++++++++++++++++++++++++++++++++++++++++-----
 3 files changed, 57 insertions(+), 10 deletions(-)
 rename pip.go => python.go (59%)

diff --git a/README.md b/README.md
index 5e749b6..748fe0c 100644
--- a/README.md
+++ b/README.md
@@ -1,7 +1,7 @@
 # Confused
 
 A tool for checking for lingering free namespaces for private package names referenced in dependency configuration
-for Python (pypi) `requirements.txt`, JavaScript (npm) `package.json`, PHP (composer) `composer.json` or MVN (maven) `pom.xml`.
+for Python (pypi) `requirements.txt` and `Pipfile`, JavaScript (npm) `package.json`, PHP (composer) `composer.json` or MVN (maven) `pom.xml`.
 
 ## What is this all about?
 
@@ -46,7 +46,7 @@ Usage:
 
 Usage of ./confused:
   -l string
-        Package repository system. Possible values: "pip", "npm", "composer", "mvn" (default "npm")
+        Package repository system. Possible values: "pip", "pipenv", "npm", "composer", "mvn" (default "npm")
   -s string
         Comma-separated list of known-secure namespaces. Supports wildcards
   -v    Verbose output
@@ -64,6 +64,14 @@ Issues found, the following packages are not available in public package reposit
 
 ```
 
+```
+./confused -l pipenv Pipfile
+
+Issues found, the following packages are not available in public package repositories:
+ [!] internal_package1
+
+```
+
 ### JavaScript (npm)
 ```
 ./confused -l npm package.json
diff --git a/main.go b/main.go
index 4a39acc..a141634 100644
--- a/main.go
+++ b/main.go
@@ -21,7 +21,7 @@ func main() {
 	verbose := false
 	filename := ""
 	safespaces := ""
-	flag.StringVar(&lang, "l", "npm", "Package repository system. Possible values: \"pip\", \"npm\", \"composer\", \"mvn\"")
+	flag.StringVar(&lang, "l", "npm", "Package repository system. Possible values: \"pip\", \"pipenv\", \"npm\", \"composer\", \"mvn\"")
 	flag.StringVar(&safespaces, "s", "", "Comma-separated list of known-secure namespaces. Supports wildcards")
 	flag.BoolVar(&verbose, "v", false, "Verbose output")
 	flag.Parse()
@@ -34,8 +34,8 @@ func main() {
 	}
 
 	filename = flag.Args()[0]
-	if lang == "pip" {
-		resolver = NewPythonLookup(verbose)
+	if lang == "pip" || lang == "pipenv" {
+		resolver = NewPythonLookup(verbose, lang)
 	} else if lang == "npm" {
 		resolver = NewNPMLookup(verbose)
 	} else if lang == "composer" {
diff --git a/pip.go b/python.go
similarity index 59%
rename from pip.go
rename to python.go
index 129bf65..6456c95 100644
--- a/pip.go
+++ b/python.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"fmt"
+	"github.com/pelletier/go-toml"
 	"io/ioutil"
 	"net/http"
 	"strings"
@@ -9,19 +10,34 @@ import (
 
 // PythonLookup represents a collection of python packages to be tested for dependency confusion.
 type PythonLookup struct {
-	Packages []string
-	Verbose  bool
+	Packages       []string
+	Verbose        bool
+	PackageManager string
 }
 
 // NewPythonLookup constructs a `PythonLookup` struct and returns it
-func NewPythonLookup(verbose bool) PackageResolver {
-	return &PythonLookup{Packages: []string{}, Verbose: verbose}
+func NewPythonLookup(verbose bool, packageManager string) PackageResolver {
+	return &PythonLookup{Packages: []string{}, Verbose: verbose, PackageManager: packageManager}
 }
 
-// ReadPackagesFromFile reads package information from a python `requirements.txt` file
+// ReadPackagesFromFile chooses a file parser based on the user-supplied python package manager.
 //
 // Returns any errors encountered
 func (p *PythonLookup) ReadPackagesFromFile(filename string) error {
+	switch p.PackageManager {
+	case "pip":
+		return p.ReadPackagesFromRequirementsTxt(filename)
+	case "pipenv":
+		return p.ReadPackagesFromPipfile(filename)
+	default:
+		return fmt.Errorf("Python package manager not implemented: %s", p.PackageManager)
+	}
+}
+
+// ReadPackagesFromRequirementsTxt reads package information from a python `requirements.txt`.
+//
+// Returns any errors encountered
+func (p *PythonLookup) ReadPackagesFromRequirementsTxt(filename string) error {
 	rawfile, err := ioutil.ReadFile(filename)
 	if err != nil {
 		return err
@@ -50,6 +66,29 @@ func (p *PythonLookup) ReadPackagesFromFile(filename string) error {
 	return nil
 }
 
+// ReadPackagesFromPipfile reads package information from a python `Pipfile`.
+//
+// Returns any errors encountered
+func (p *PythonLookup) ReadPackagesFromPipfile(filename string) error {
+	rawfile, err := ioutil.ReadFile(filename)
+	if err != nil {
+		return err
+	}
+	config, err := toml.Load(string(rawfile))
+	if err != nil {
+		return err
+	}
+	packages := config.Get("packages")
+	if packages != nil {
+		p.Packages = append(p.Packages, packages.(*toml.Tree).Keys()...)
+	}
+	dev_packages := config.Get("dev-packages")
+	if dev_packages != nil {
+		p.Packages = append(p.Packages, dev_packages.(*toml.Tree).Keys()...)
+	}
+	return nil
+}
+
 // PackagesNotInPublic determines if a python package does not exist in the pypi package repository.
 //
 // Returns a slice of strings with any python packages not in the pypi package repository